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

A node-based canvas (gr.Workflow) exposing the diffusers QwenImage21Pipeline:

  - enhance_prompt_t2i: prompt -> rewritten prompt   (Qwen/Qwen-Image-2.1-PE-T2I)
  - text_to_image:      prompt -> image              (Qwen/Qwen-Image-2.1)
  - enhance_prompt_i2i: image + instruction -> rewritten instruction
                                                     (Qwen/Qwen-Image-2.1-PE-I2I)
  - edit_image:         condition image + instruction -> edited image

Image-input requests pass through an NCII prompt classifier
(hfmlsoc/ncii-guard-v02) before any rewriting or editing runs.

All models are placed on `cuda` at module level, as ZeroGPU requires —
the diffusion pipeline and both 9B prompt-rewriting models stay resident
(~66GB of weights, fits the 96GB xlarge card), so enhance calls pay no
transfer time. The 270M guard classifier runs on CPU.
"""

import base64
import io
import json
import os
import urllib.parse
import urllib.request

import gradio as gr
import spaces
import torch
from gradio_client import utils as client_utils
from gradio.utils import get_upload_folder
from huggingface_hub import hf_hub_download
from PIL import Image

from diffusers import QwenImage21Pipeline

MODEL_ID = os.environ.get("QWEN_IMAGE_MODEL", "Qwen/Qwen-Image-2.1")
PE_T2I_ID = "Qwen/Qwen-Image-2.1-PE-T2I"
PE_I2I_ID = "Qwen/Qwen-Image-2.1-PE-I2I"
GUARD_ID = "hfmlsoc/ncii-guard-v02"
GUARD_THRESHOLD = 0.5

# All model repos are public. If HF_TOKEN is set as a Space secret it is used
# implicitly for Hub requests (higher rate limits), but it is not required.
pipe = QwenImage21Pipeline.from_pretrained(MODEL_ID, dtype=torch.bfloat16)
pipe.to("cuda")

# --- Prompt rewriting models (resident in CPU RAM, moved to GPU per call) ---
from transformers import (
    AutoModelForCausalLM,
    AutoModelForImageTextToText,
    AutoProcessor,
    AutoTokenizer,
)

pe_t2i_tokenizer = AutoTokenizer.from_pretrained(PE_T2I_ID)
pe_t2i = AutoModelForCausalLM.from_pretrained(PE_T2I_ID, dtype=torch.bfloat16).to("cuda").eval()
pe_t2i_system = open(hf_hub_download(PE_T2I_ID, "system_prompt.txt")).read().strip()

pe_i2i_processor = AutoProcessor.from_pretrained(PE_I2I_ID)
pe_i2i = AutoModelForImageTextToText.from_pretrained(PE_I2I_ID, dtype=torch.bfloat16).to("cuda").eval()
pe_i2i_system = open(hf_hub_download(PE_I2I_ID, "system_prompt.txt")).read().strip()

# --- NCII guard for image-input requests (runs on CPU) ---
from transformers import AutoModelForSequenceClassification

guard_tokenizer = AutoTokenizer.from_pretrained(GUARD_ID)  # carries the normalizer
guard = AutoModelForSequenceClassification.from_pretrained(GUARD_ID).eval()


def _check_prompt_guard(prompt: str) -> None:
    """Reject image-editing prompts the NCII classifier flags. The error is
    deliberately generic and does not say which classifier fired."""
    batch = guard_tokenizer(
        [prompt], truncation=True, max_length=256, padding=True, return_tensors="pt"
    )
    with torch.no_grad():
        prob = torch.softmax(guard(**batch).logits.float(), dim=-1)[0, 1].item()
    if prob >= GUARD_THRESHOLD:
        raise gr.Error("prompt invalid based on our classifiers, try again")


def _to_pil(image) -> Image.Image:
    """Accept whatever the canvas hands a bound function for an image port:
    a PIL image, a local path, a /gradio_api/file= reference, an http(s) or
    data: URL, or a file dict carrying any of those. Mirrors _file_ref in
    gradio.workflow — canvas file values carry only `url`, no `path`."""
    if isinstance(image, Image.Image):
        return image
    if isinstance(image, dict):
        image = image.get("path") or image.get("url") or ""
    if not isinstance(image, str) or not image:
        raise ValueError(f"Unsupported image input: {type(image)!r}")
    if image.startswith("/gradio_api/file="):
        image = urllib.parse.unquote(image.removeprefix("/gradio_api/file="))
    if image.startswith("data:"):
        return Image.open(io.BytesIO(base64.b64decode(image.split(",", 1)[1])))
    if image.startswith(("http://", "https://")):
        return Image.open(io.BytesIO(urllib.request.urlopen(image).read()))
    return Image.open(image)


def _save(image: Image.Image) -> dict:
    # Mirror gradio.workflow._save_tmp: the canvas renders media values only
    # from {path, url, is_file} dicts whose url is a /gradio_api/file= link,
    # and the file must live under the upload folder to be servable.
    directory = get_upload_folder()
    os.makedirs(directory, exist_ok=True)
    path = os.path.join(directory, f"workflow_{os.urandom(8).hex()}.png")
    image.save(path)
    url = f"/gradio_api/file={client_utils.encode_file_path(path)}"
    return {"path": path, "url": url, "is_file": True}


def _steps(value, default: int = 28) -> 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)


def _parse_rewrite(gen: str, fallback: str) -> str:
    """Split the PE model's <think> block from its JSON answer and return the
    rewritten prompt. Fall back to the original prompt if parsing fails."""
    _, _, answer = gen.partition("</think>")
    try:
        return json.loads(answer.strip()).get("rewritten_prompt") or fallback
    except (json.JSONDecodeError, AttributeError):
        return fallback


def _enhance_duration_t2i(prompt: str, enhance: bool = False) -> int:
    # Dynamic duration (per ZeroGPU docs): when the Enhance checkbox is off
    # the function returns immediately, so only reserve a couple seconds of
    # GPU time; the full budget is only requested when enhancement runs.
    return 90 if enhance else 2


@spaces.GPU(size="xlarge", duration=_enhance_duration_t2i)
def enhance_prompt_t2i(prompt: str, enhance: bool = False) -> str:
    """Rewrite a brief text-to-image request into a detailed English prompt
    with Qwen-Image-2.1-PE-T2I. Passes the prompt through unchanged unless
    the Enhance checkbox is on (enhancement is opt-in to save quota)."""
    if not enhance:
        return prompt
    text = pe_t2i_tokenizer.apply_chat_template(
        [{"role": "system", "content": pe_t2i_system},
         {"role": "user", "content": prompt}],
        tokenize=False, add_generation_prompt=True, enable_thinking=True,
    )
    inputs = pe_t2i_tokenizer(text, return_tensors="pt").to("cuda")
    with torch.no_grad():
        out = pe_t2i.generate(
            **inputs, max_new_tokens=1024,
            do_sample=True, temperature=1.0, top_p=0.95, top_k=20,
        )
    gen = pe_t2i_tokenizer.decode(
        out[0, inputs["input_ids"].shape[1]:], skip_special_tokens=True
    )
    return _parse_rewrite(gen, prompt)


def _enhance_duration_i2i(image, instruction: str, enhance: bool = False) -> int:
    return 90 if enhance else 2


@spaces.GPU(size="xlarge", duration=_enhance_duration_i2i)
def enhance_prompt_i2i(image, instruction: str, enhance: bool = False) -> str:
    """Rewrite an image-editing instruction against the condition image with
    Qwen-Image-2.1-PE-I2I. Passes the instruction through unchanged unless
    the Enhance checkbox is on (enhancement is opt-in to save quota)."""
    _check_prompt_guard(instruction)
    if not enhance:
        return instruction
    pil_image = _to_pil(image).convert("RGB")
    messages = [
        {"role": "system", "content": [{"type": "text", "text": pe_i2i_system}]},
        {"role": "user", "content": [
            {"type": "image", "image": pil_image},
            {"type": "text", "text": instruction},
        ]},
    ]
    inputs = pe_i2i_processor.apply_chat_template(
        messages, add_generation_prompt=True, tokenize=True,
        return_dict=True, return_tensors="pt", enable_thinking=True,
    ).to("cuda")
    with torch.no_grad():
        out = pe_i2i.generate(
            **inputs, max_new_tokens=1024,
            do_sample=True, temperature=1.0, top_p=0.95, top_k=20,
        )
    gen = pe_i2i_processor.tokenizer.decode(
        out[0, inputs["input_ids"].shape[1]:], skip_special_tokens=True
    )
    return _parse_rewrite(gen, instruction)


# xlarge requests are billed at 2x the reserved duration, so a full canvas
# Run (T2I + Edit + 2 enhance passthroughs) reserves 2*25*2 + 2*2*2 = 108s,
# fitting the 120s/day anonymous quota. 28 steps runs in ~15s on this GPU.
@spaces.GPU(size="xlarge", duration=25)
def text_to_image(prompt: str, steps: int = 28) -> dict:
    """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(size="xlarge", duration=25)
def edit_image(image, instruction: str, steps: int = 28) -> dict:
    """Edit a condition image following an instruction (image-conditioned
    generation with Qwen-Image 2.1)."""
    _check_prompt_guard(instruction)
    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={
        "enhance_prompt_t2i": enhance_prompt_t2i,
        "enhance_prompt_i2i": enhance_prompt_i2i,
        "text_to_image": text_to_image,
        "edit_image": edit_image,
    },
)

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