from pathlib import Path

import replicate


def flux_schnell(prompt: str, file_path: Path):
    output = replicate.run(
        "black-forest-labs/flux-schnell",
        input={
            "prompt": prompt,
            "go_fast": True,
            "megapixels": "0.25",
            "num_outputs": 1,
            "aspect_ratio": "1:1",
            "output_format": "webp",
            "output_quality": 80,
            "num_inference_steps": 4,
        },
    )
    with open(file_path, "wb") as fp:
        fp.write(output[0].read())


def nkodeicon_v1(prompt: str, file_path: Path):
    output = replicate.run(
        "donovankelly/nkodeicon-v1:761f0a14dbef2a9cd4d65fbb19f613209f7e3eb151c42ec9a5ee775d30e83e9b",
        input={
            "model": "dev",
            "prompt": f"TOK {prompt} on a white background",
            "go_fast": True,
            "lora_scale": 1.5,
            "megapixels": "0.25",
            "num_outputs": 1,
            "aspect_ratio": "1:1",
            "output_format": "webp",
            "guidance_scale": 3.5,
            "output_quality": 100,
            "prompt_strength": 0.8,
            "extra_lora_scale": 1.5,
            "num_inference_steps": 28,
        },
    )
    with open(file_path, "wb") as fp:
        fp.write(output[0].read())


def nkodeicon_v3(prompt: str, file_path: Path):
    output = replicate.run(
        "donovankelly/nkodeicon-v3:12fa3b15dd48dc3081642572f2cb7fd18652086f19133dddd4b666f4c7455752",
        input={
            "model": "dev",
            "prompt": f"ICON colorful {prompt} on a white background",
            "go_fast": True,
            "lora_scale": 1,
            "megapixels": "0.25",
            "num_outputs": 1,
            "aspect_ratio": "1:1",
            "output_format": "webp",
            "guidance_scale": 3,
            "output_quality": 80,
            "prompt_strength": 0.8,
            "extra_lora_scale": 1,
            "num_inference_steps": 50,
        },
    )
    with open(file_path, "wb") as fp:
        fp.write(output[0].read())