"""gr.Workflow app for PrunaAI/Pruna-Qwen-Image-2.1 (ZeroGPU).

Nodes on the canvas:
  - "Generate Image": text-to-image with the 5- or 8-step LoRA adapter
  - "Edit Image": image editing with an input image + prompt
"""

import os

import torch
import spaces
import gradio as gr
from gradio.utils import get_upload_folder
from gradio_client import utils as client_utils
from PIL import Image
from diffusers import FlowMatchEulerDiscreteScheduler, QwenImage21Pipeline

SIGMAS = {
    5: [1.0, 0.94, 6 / 7, 2 / 3, 0.4],
    8: [1.0, 14 / 15, 6 / 7, 10 / 13, 2 / 3, 6 / 11, 0.4, 2 / 9],
}

# Load the base pipeline once at module level (ZeroGPU CUDA emulation makes
# this a no-op until a @spaces.GPU function actually runs).
pipe = QwenImage21Pipeline.from_pretrained(
    "Qwen/Qwen-Image-2.1", torch_dtype=torch.bfloat16
).to("cuda")
pipe.scheduler = FlowMatchEulerDiscreteScheduler.from_config(
    pipe.scheduler.config,
    use_dynamic_shifting=False,
    shift=1.0,
    shift_terminal=None,
)

_loaded_steps = None


def _ensure_adapter(steps: int):
    """Load exactly one Pruna LoRA adapter matching the requested step count."""
    global _loaded_steps
    steps = 5 if steps is not None and int(steps) == 5 else 8
    if _loaded_steps != steps:
        pipe.unload_lora_weights()
        pipe.load_lora_weights(
            "PrunaAI/Pruna-Qwen-Image-2.1",
            weight_name=f"p_qwen_image_2.1_{steps}step_v0.1.safetensors",
        )
        _loaded_steps = steps
    return steps


def _seed(seed) -> int:
    try:
        return int(seed)
    except (TypeError, ValueError):
        return 42


def _coerce_image(image):
    """Accept a PIL image, a filepath, or a Gradio/Workflow file-ref dict."""
    if image is None or image == "":
        return None
    if isinstance(image, Image.Image):
        return image.convert("RGB")
    if isinstance(image, dict):
        image = image.get("path") or image.get("url") or ""
    return Image.open(str(image)).convert("RGB")


def _to_file(img: Image.Image) -> dict:
    """Save a PIL image into Gradio's upload folder and return a file-ref dict
    (same shape gr.Workflow uses internally for media node outputs)."""
    directory = get_upload_folder()
    os.makedirs(directory, exist_ok=True)
    path = os.path.join(directory, f"pruna_{os.urandom(8).hex()}.png")
    img.save(path)
    url = f"/gradio_api/file={client_utils.encode_file_path(path)}"
    return {"path": path, "url": url, "is_file": True}


@spaces.GPU(duration=120)
def generate(prompt: str, steps: int = 8, seed: int = 42) -> dict:
    """Text-to-image with Pruna-Qwen-Image-2.1 (5 or 8 steps, no CFG)."""
    steps = _ensure_adapter(steps)
    out = pipe(
        prompt=prompt,
        width=1024,
        height=1024,
        generator=torch.Generator("cuda").manual_seed(_seed(seed)),
        num_inference_steps=steps,
        sigmas=SIGMAS[steps],
        true_cfg_scale=1.0,
        use_kv_cache=True,
    ).images[0]
    return _to_file(out)


@spaces.GPU(duration=120)
def edit(image, prompt: str, steps: int = 8, seed: int = 42) -> dict:
    """Edit an image with Pruna-Qwen-Image-2.1 (5 or 8 steps, no CFG)."""
    steps = _ensure_adapter(steps)
    out = pipe(
        prompt=prompt,
        image=_coerce_image(image),
        generator=torch.Generator("cuda").manual_seed(_seed(seed)),
        num_inference_steps=steps,
        sigmas=SIGMAS[steps],
        true_cfg_scale=1.0,
        use_kv_cache=True,
    ).images[0]
    return _to_file(out)


demo = gr.Workflow(
    graph="workflow.json",
    bind={"Generate Image": generate, "Edit Image": edit},
)

if __name__ == "__main__":
    demo.launch()