import yaml
import textwrap

IMAGE = "quay.io/lukemarsden/lora:v0.0.4"
MODEL_NAME = "runwayml/stable-diffusion-v1-5"

# Don't want to make this variable since the job is fixed price
NUM_IMAGES = 10

def _lora_training(params: str):
    if params.startswith("{"):
        params = yaml.safe_load(params)
    else:
        raise Exception("Please set params to yaml like {seed: 42, 'images_cid': 'Qm...'} where images_cid contains an images.zip with training images in")

    seed = params.get("seed", 42)
    images_cid = params["images_cid"]

    return {
        "APIVersion": "V1beta1",
        "Metadata": {
            "CreatedAt": "0001-01-01T00:00:00Z",
            "Requester": {}
        },
        "Spec": {
            "Deal": {
                "Concurrency": 1
            },
            "Docker": {
                "EnvironmentVariables": [
                    f"RANDOM_SEED={seed}",
                ],
                "Entrypoint": [
                    "bash", "-c", f'(cd /input && unzip images.zip && rm images.zip) && lora_pti --pretrained_model_name_or_path={MODEL_NAME} --instance_data_dir=/input --output_dir=/output --train_text_encoder --resolution=512 --train_batch_size=1 --gradient_accumulation_steps=4 --scale_lr --learning_rate_unet=1e-4 --learning_rate_text=1e-5 --learning_rate_ti=5e-4 --color_jitter --lr_scheduler="linear" --lr_warmup_steps=0 --placeholder_tokens="<s1>|<s2>" --use_template="style" --save_steps=100 --max_train_steps_ti=1000 --max_train_steps_tuning=1000 --perform_inversion=True --clip_ti_decay --weight_decay_ti=0.000 --weight_decay_lora=0.001 --continue_inversion --continue_inversion_lr=1e-4 --device="cuda:0" --lora_rank=1'
                ],
                "Image": IMAGE,
            },
            "Engine": "Docker",
            "Language": {
                "JobContext": {}
            },
            "Network": {
                "Type": "None"
            },
            "PublisherSpec": {
                "Type": "Estuary"
            },
            "Resources": {
                "GPU": "1"
            },
            "Timeout": 1800,
            "Verifier": "Noop",
            "Wasm": {
                "EntryModule": {}
            },
            "inputs": [
                {
                    "CID": images_cid,
                    "Name": "lora_input",
                    "StorageSource": "IPFS",
                    "path": "/input",

                },
            ],
            "outputs": [
                {
                    "Name": "output",
                    "StorageSource": "IPFS",
                    "path": "/output"
                }
            ]
        }
    }

def _lora_inference(params: str):
    if params.startswith("{"):
        params = yaml.safe_load(params)
    else:
        raise Exception("Please set params to yaml like {seed: 42, 'lora_cid': 'Qm...', "+
                        "prompt: 'an astronaut in the style of <s1><s2>'} "+
                        "where lora_cid is the output cid of the above step")

    # TODO add a default we pin
    lora_cid = params["lora_cid"]
    seed = params.get("seed", 42)
    prompt = params.get("prompt", "question mark floating in space")
    finetune_weighting = params.get("finetune_weighting", 0.5)

    return {
        "APIVersion": "V1beta1",
        "Metadata": {
            "CreatedAt": "0001-01-01T00:00:00Z",
            "Requester": {}
        },
        "Spec": {
            "Deal": {
                "Concurrency": 1
            },
            "Docker": {
                "EnvironmentVariables": [
                    f"PROMPT={prompt}",
                    f"RANDOM_SEED={seed}",
                    f"FINETUNE_WEIGHTING={finetune_weighting}",
                    "HF_HUB_OFFLINE=1",
                ],
                "Entrypoint": [
                    'python3',
                    '-c',
                    # dedent
                    textwrap.dedent(f"""
                        from diffusers import StableDiffusionPipeline, EulerAncestralDiscreteScheduler
                        import os
                        import torch
                        from lora_diffusion import tune_lora_scale, patch_pipe

                        model_id = "{MODEL_NAME}"

                        pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16).to(
                            "cuda"
                        )
                        pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)

                        prompt = os.getenv("PROMPT")
                        seed = int(os.getenv("RANDOM_SEED"))
                        torch.manual_seed(seed)

                        os.system("find /input")
                        patch_pipe(
                            pipe,
                            "/input/output/final_lora.safetensors",
                            patch_text=True,
                            patch_ti=True,
                            patch_unet=True,
                        )

                        coeff = float(os.getenv("FINETUNE_WEIGHTING", 0.5))
                        tune_lora_scale(pipe.unet, coeff)
                        tune_lora_scale(pipe.text_encoder, coeff)

                        image = pipe(prompt, num_inference_steps=50, guidance_scale=7).images[0]
                        image.save(f"/output/image-{{seed}}.jpg")
                        image
                        """)
                ],
                "Image": IMAGE,
            },
            "Engine": "Docker",
            "Language": {
                "JobContext": {}
            },
            "Network": {
                "Type": "None"
            },
            "PublisherSpec": {
                "Type": "Estuary"
            },
            "Resources": {
                "GPU": "1"
            },
            "Timeout": 1800,
            "Verifier": "Noop",
            "Wasm": {
                "EntryModule": {}
            },
            "inputs": [
                {
                    "CID": lora_cid,
                    "Name": "lora_input",
                    "StorageSource": "IPFS",
                    "path": "/input",

                },
            ],
            "outputs": [
                {
                    "Name": "output",
                    "StorageSource": "IPFS",
                    "path": "/output"
                },
            ]
        }
    }

import json

def _fastchat(params: str):
    cmd = ["python3", "main.py"]
    return {
        "APIVersion": "V1beta1",
        "Metadata": {
            "CreatedAt": "0001-01-01T00:00:00Z",
            "Requester": {}
        },
        "Spec": {
            "Deal": {
                "Concurrency": 1
            },
            "Docker": {
                "Entrypoint": cmd,
                "Image": "xqua/carpai-demo-repo:v0.6",
            },
            "Engine": "Docker",
            "Language": {
                "JobContext": {}
            },
            "Network": {
                "Type": "None"
            },
            "PublisherSpec": {
                "Type": "Estuary"
            },
            "Resources": {
                "GPU": "1"
            },
            "Timeout": 1800,
            "Verifier": "Noop",
            "Wasm": {
                "EntryModule": {}
            },
            "inputs": [
                {
                    "CID": params,
                    "Name": "prompt_template_input",
                    "StorageSource": "IPFS",
                    "path": "/prompt_template.json",

                },
            ],
            "outputs": [
                {
                    "Name": "output",
                    "StorageSource": "IPFS",
                    "path": "/output"
                }
            ]
        }
    }

if __name__ == "__main__":
    print(_fastchat('QmcPjQwVcJiFge3yNjVL2NoZsTQ3GBpXAZe21S2Ncg16Gt'))

import yaml

def _sdxl(params: str):
    if params.startswith("{"):
        params = yaml.safe_load(params)
    else:
        prompt = params
        params = {"prompt": prompt, "seed": 0}
    if not isinstance(params, dict):
        raise Exception("Please set params to a dict like {'prompt': 'astronaut riding a horse', 'seed': 42}")
    return {
        "APIVersion": "V1beta1",
        "Metadata": {
            "CreatedAt": "0001-01-01T00:00:00Z",
            "Requester": {}
        },
        "Spec": {
            "Deal": {
                "Concurrency": 1
            },
            "Docker": {
                "Entrypoint": [
                    "bash", "-c",
                    # stderr logging is nondeterministic (includes timing information)
                    "python3 inference.py 2>/dev/null",
                ],
                "Image": "quay.io/lukemarsden/sdxl:v0.9-lilypad1-v2",
                "EnvironmentVariables": [
                    f"PROMPT={params.get('prompt', 'question mark floating in space')}",
                    f"RANDOM_SEED={params.get('seed', 0)}",
                    f"OUTPUT_DIR=/outputs/",
                    "HF_HUB_OFFLINE=1",
                ]
            },
            "Engine": "Docker",
            "Language": {
                "JobContext": {}
            },
            "Network": {
                "Type": "None"
            },
            "PublisherSpec": {
                "Type": "Estuary"
            },
            "Resources": {
                "GPU": "1"
            },
            "Timeout": 1800,
            "Verifier": "Noop",
            "Wasm": {
                "EntryModule": {}
            },
            "outputs": [
                {
                    "Name": "outputs",
                    "StorageSource": "IPFS",
                    "path": "/outputs"
                }
            ]
        }
    }

if __name__ == "__main__":
    print(_sdxl("{prompt: hello, seed: 99}"))
    print(_sdxl("{prompt: 'hello world', seed: 99}"))

// SPDX-License-Identifier: GPLv3
pragma solidity ^0.8.6;

interface ModicumContract {
  function runModuleWithDefaultMediators(string calldata name, string calldata params) external payable returns (uint256);
}

// Payment is 2 lilETH for all jobs currently
// got to testnet.lilypadnetwork.org to fund your wallet
contract ExampleClient {
  address public _contractAddress;
  ModicumContract remoteContractInstance;

  uint256 public lilypadFee = 2;

  struct Result {
      uint256 jobID;
      string cid;
      string httpString;
  }

  Result[] public results;

  event ReceivedJobResults(uint256 jobID, string cid);

  // The Modicum contract address is found here: https://github.com/bacalhau-project/lilypad-modicum/blob/main/latest.txt
  // Current: 0x422F325AA109A3038BDCb7B03Dd0331A4aC2cD1a
  constructor (address contractAddress) {
    require(contractAddress != address(0), "NaiveExamplesClient: contract cannot be zero address");
    _contractAddress = contractAddress;
    remoteContractInstance = ModicumContract(contractAddress);
  }

  function setLilypadFee(uint256 _fee) public {
      lilypadFee = _fee;
  }
  
  function runCowsay(string memory sayWhat) public payable returns (uint256) {
    require(msg.value == lilypadFee * 1 ether, string(abi.encodePacked("Payment of 2 Ether is required")));
    return runModule("cowsay:v0.0.1", sayWhat);
  }

  function runStablediffusion(string memory prompt) public payable returns (uint256) {
    require(msg.value == lilypadFee * 1 ether, string(abi.encodePacked("Payment of 2 Ether is required")));
    return runModule("stable_diffusion:v0.0.1", prompt);
  }

  function runSDXL(string memory prompt) public payable returns (uint256) {
    require(msg.value == lilypadFee * 1 ether, string(abi.encodePacked("Payment of 2 Ether is required")));
    return runModule("sdxl:v0.9-lilypad1", prompt);
  }

 function runModule(string memory name, string memory params) public payable returns (uint256) {
    require(msg.value == lilypadFee * 1 ether, string(abi.encodePacked("Payment of 2 Ether is required")));
    return remoteContractInstance.runModuleWithDefaultMediators{value: msg.value}(name, params);
  }

  // Implemented Modicum interface. This saves results to a results array
  function receiveJobResults(uint256 _jobID, string calldata _cid) public {
    Result memory jobResult = Result({
        jobID: _jobID,
        cid: _cid,
        httpString: string(abi.encodePacked("https://ipfs.io/ipfs/", _cid))
    });
    results.push(jobResult);

    emit ReceivedJobResults(_jobID, _cid);
  }

  function fetchAllResults() public view returns (Result[] memory) {
    return results;
  }
}

import yaml

def _duckdb(params: str):
    if params.startswith("{"):
        params = yaml.safe_load(params)
    else:
        raise Exception("Please set params to yaml like {'query': 'SELECT...', 'inputs_cid': 'Qm...'}")

    return {
        "APIVersion": "V1beta1",
        "Metadata": {
            "CreatedAt": "0001-01-01T00:00:00Z",
            "Requester": {}
        },
        "Spec": {
            "Engine": "Docker",
            "Verifier": "Noop",
            "Publisher": "Estuary",
            "PublisherSpec": {
                "Type": "Estuary"
            },
            "Docker": {
                "Image": "31z4/bacalhau-duckdb:0.0.1",
                "Entrypoint": [
                    "./duckdb",
                    "-init",
                    "/init.sql",
                    "-echo",
                    "-s",
                    params["query"]
                ],
            },
            "Language": {
                "JobContext": {}
            },
            "Wasm": {
                "EntryModule": {}
            },
            "Resources": {
                "GPU": ""
            },
            "Network": {
                "Type": "None"
            },
            "Timeout": 1800,
            "inputs": [
                {
                    "StorageSource": "IPFS",
                    "Name": "inputs",
                    "CID": params["inputs_cid"],
                    "path": "/inputs",
                }
            ],
            "outputs": [
                {
                    "StorageSource": "IPFS",
                    "Name": "outputs",
                    "path": "/outputs"
                }
            ],
            "Deal": {
                "Concurrency": 1
            },
        },
    }