"""Qwen-Image 2.1 — Gradio Workflow on ZeroGPU.

A node-based canvas (gr.Workflow) exposing the diffusers QwenImage21Pipeline:
  - text_to_image: prompt -> image
  - edit_image:    condition image + instruction -> edited image (chained after
                   text_to_image, or fed from an uploaded image reference)

The pipeline is placed on `cuda` at module level: outside @spaces.GPU functions
PyTorch runs in CUDA emulation mode, and a real ZeroGPU is attached only while
a decorated function executes.
"""

import os
import tempfile

import gradio as gr
import spaces
import torch
from PIL import Image

from diffusers import QwenImage21Pipeline

MODEL_ID = os.environ.get("QWEN_IMAGE_MODEL", "Qwen/Qwen-Image-2.1")

# The model repo is private/gated. Per the Hub auth docs, the HF_TOKEN
# environment variable is used implicitly for all Hub requests and takes
# priority over any stored token — so set HF_TOKEN as a Space secret
# (Settings -> Secrets) with a token whose account has access to the repo.
if not os.environ.get("HF_TOKEN"):
    raise RuntimeError(
        "HF_TOKEN is not set. Add it as a Space secret (Settings -> Secrets) "
        f"with read access to {MODEL_ID}."
    )

pipe = QwenImage21Pipeline.from_pretrained(MODEL_ID, dtype=torch.bfloat16)
pipe.to("cuda")


def _to_pil(image) -> Image.Image:
    """Accept whatever the canvas hands a bound function for an image port:
    a PIL image, a file path, or Gradio's {path, url, is_file} file dict."""
    if isinstance(image, Image.Image):
        return image
    if isinstance(image, dict):
        image = image.get("path") or image.get("url")
    if isinstance(image, str):
        return Image.open(image)
    raise ValueError(f"Unsupported image input: {type(image)!r}")


def _save(image: Image.Image) -> str:
    path = tempfile.NamedTemporaryFile(suffix=".png", delete=False).name
    image.save(path)
    return path


def _steps(value, default: int = 40) -> int:
    # Unconnected number ports arrive as None; function defaults are not
    # applied because the workflow executor passes arguments positionally.
    return default if value is None else int(value)


@spaces.GPU(duration=120)
def text_to_image(prompt: str, steps: int = 40) -> str:
    """Generate an image from a text prompt with Qwen-Image 2.1."""
    image = pipe(prompt, num_inference_steps=_steps(steps)).images[0]
    return _save(image)


@spaces.GPU(duration=120)
def edit_image(image, instruction: str, steps: int = 40) -> str:
    """Edit a condition image following an instruction (image-conditioned
    generation with Qwen-Image 2.1)."""
    edited = pipe(instruction, image=_to_pil(image), num_inference_steps=_steps(steps)).images[0]
    return _save(edited)


demo = gr.Workflow(
    graph=os.path.join(os.path.dirname(__file__), "workflow.json"),
    bind={"text_to_image": text_to_image, "edit_image": edit_image},
)

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