#!/usr/bin/env python3
"""
Simple Gradio UI for Qwen-Image-2.1, driving ComfyUI's own API.
Generated by Model Lab. Pairs with setup_qwen_image_2_1_comfyui.py.

No node graph, no ComfyUI account, nothing to install by hand beyond this
one file: three tabs (Text to Image, Image Edit, Remove Background), each
just a prompt box and/or an image upload and a Generate button. Every tab
sends a plain ComfyUI /prompt API call to whichever ComfyUI instance you
point it at (defaults to localhost:8188, i.e. this same pod) and polls
/history until the result is ready, then displays it.

Normally you don't run this file directly - setup_qwen_image_2_1_comfyui.py's
--serve step finds it automatically if it's sitting in the same folder and
launches it on port 8889 (already in the guide's standard "Expose HTTP
ports" list, so nothing extra to expose). Running it by hand only matters if
you want a different port or you're driving a ComfyUI instance elsewhere:

    pip install gradio requests pillow
    python qwen_image_2_1_gradio.py --port 8889

Whichever port you use, expose it the same way you exposed 8188 (RunPod: add
an HTTP port mapping), then open that URL. Add ?__theme=light/dark if you
want a specific theme; Gradio defaults to the browser's preference.

Model filenames match setup_qwen_image_2_1_comfyui.py's int8 default
(qwen_image_2.1_int8_convrot.safetensors / qwen3vl_8b_w4a8.safetensors /
qwen_image_2.1_vae_bf16.safetensors). If you ran the setup script with
--full or --gguf, change MODEL_FILES below to match, or pass
--unet/--clip/--vae on the command line - the setup script's own --serve
step does this automatically when it launches this file for you.

A .gguf unet filename automatically switches the loader node from core
ComfyUI's UNETLoader to city96/ComfyUI-GGUF's UnetLoaderGGUF (see
_unet_loader_node() below) - the setup script clones that custom node
whenever --gguf is used, so it's already present by the time this launches.
"""

import argparse
import base64
import io
import json
import os
import random
import time
import uuid
from typing import Optional

import gradio as gr
import requests
from PIL import Image

MODEL_FILES = {
    "unet": "qwen_image_2.1_int8_convrot.safetensors",
    "clip": "qwen3vl_8b_w4a8.safetensors",
    "vae": "qwen_image_2.1_vae_bf16.safetensors",
}

BG_REMOVAL_PROMPT = "Remove the background entirely. Keep only the main subject, unchanged, with clean edges and a fully transparent background."


def _client(base_url: str) -> "ComfyClient":
    return ComfyClient(base_url)


class ComfyClient:
    """Minimal wrapper around ComfyUI's /prompt + /history + /view API.

    ComfyUI's REST API takes graphs in "API format" (a flat
    {node_id: {class_type, inputs}} dict) - not the "UI format" the graph
    editor saves (nodes/links/groups), which is a different shape entirely.
    The graphs below are written directly in API format rather than
    converted from the official templates, since hand-writing a small graph
    is far more reliable than round-tripping through ComfyUI's own
    save-as-API-format export for a handful of fixed node types.
    """

    def __init__(self, base_url: str):
        self.base_url = base_url.rstrip("/")
        self.client_id = str(uuid.uuid4())

    def upload_image(self, pil_image: Image.Image) -> str:
        """Every call MUST get its own filename. A prior version of this
        method sent every upload as the literal name "upload.png" with
        overwrite=true - fine for a single-image submit, but for
        multi-reference editing the second upload silently overwrote the
        first ON DISK inside ComfyUI's input folder before generation ever
        ran, so both LoadImage nodes ended up pointing at the same one
        file. The model never saw two distinct references; it just rendered
        one photo twice with no error anywhere in the pipeline. A random
        per-upload filename (still overwrite=true, in case of a genuine
        retry with an identical name) makes concurrent/sequential uploads
        from one Gradio submission independent of each other."""
        buf = io.BytesIO()
        pil_image.save(buf, format="PNG")
        buf.seek(0)
        filename = f"gradio_upload_{uuid.uuid4().hex}.png"
        resp = requests.post(
            f"{self.base_url}/upload/image",
            files={"image": (filename, buf, "image/png")},
            data={"overwrite": "true"},
            timeout=60,
        )
        resp.raise_for_status()
        return resp.json()["name"]

    def queue_prompt(self, graph: dict) -> str:
        resp = requests.post(
            f"{self.base_url}/prompt",
            json={"prompt": graph, "client_id": self.client_id},
            timeout=30,
        )
        if resp.status_code != 200:
            raise RuntimeError(f"ComfyUI rejected the graph: {resp.status_code} {resp.text[:500]}")
        return resp.json()["prompt_id"]

    def wait_for_result(self, prompt_id: str, save_node_id: str, timeout_s: int = 300) -> Image.Image:
        deadline = time.time() + timeout_s
        while time.time() < deadline:
            resp = requests.get(f"{self.base_url}/history/{prompt_id}", timeout=30)
            resp.raise_for_status()
            history = resp.json()
            entry = history.get(prompt_id)
            if entry is not None:
                status = entry.get("status", {})
                if status.get("status_str") == "error" or (status.get("completed") is False and status.get("messages")):
                    for msg_type, msg_data in status.get("messages", []):
                        if msg_type == "execution_error":
                            raise RuntimeError(f"ComfyUI execution error: {msg_data.get('exception_message', msg_data)}")
                outputs = entry.get("outputs", {})
                node_out = outputs.get(save_node_id)
                if node_out and node_out.get("images"):
                    img_info = node_out["images"][0]
                    return self._fetch_image(img_info["filename"], img_info.get("subfolder", ""), img_info.get("type", "output"))
            time.sleep(1.5)
        raise TimeoutError(f"Generation did not finish within {timeout_s}s. Check ComfyUI's own console for stuck/errored jobs.")

    def _fetch_image(self, filename: str, subfolder: str, folder_type: str) -> Image.Image:
        resp = requests.get(
            f"{self.base_url}/view",
            params={"filename": filename, "subfolder": subfolder, "type": folder_type},
            timeout=60,
        )
        resp.raise_for_status()
        return Image.open(io.BytesIO(resp.content))


def _unet_loader_node() -> dict:
    """A .gguf diffusion model (the --gguf setup tier) is not readable by
    core ComfyUI's UNETLoader - it needs city96/ComfyUI-GGUF's own
    UnetLoaderGGUF node instead, which the setup script clones automatically
    whenever --gguf is used. Same MODEL output either way, so nothing else
    in the graph needs to know which loader produced it."""
    if MODEL_FILES["unet"].lower().endswith(".gguf"):
        return {"class_type": "UnetLoaderGGUF", "inputs": {"unet_name": MODEL_FILES["unet"]}}
    return {"class_type": "UNETLoader", "inputs": {"unet_name": MODEL_FILES["unet"], "weight_dtype": "default"}}


def _build_t2i_graph(prompt: str, negative_prompt: str, width: int, height: int, steps: int, cfg: float, seed: int) -> dict:
    """API-format graph for the seven-stage text-to-image pipeline: load
    checkpoint -> encode prompts -> empty latent -> sample -> decode -> save.
    Node ids below are arbitrary strings (ComfyUI's API accepts any unique
    key), chosen to read clearly rather than matching the official
    template's numeric ids, which this graph does not reuse."""
    return {
        "unet_loader": _unet_loader_node(),
        "clip_loader": {"class_type": "CLIPLoader", "inputs": {"clip_name": MODEL_FILES["clip"], "type": "qwen_image", "device": "default"}},
        "vae_loader": {"class_type": "VAELoader", "inputs": {"vae_name": MODEL_FILES["vae"]}},
        "text_encode": {
            "class_type": "TextEncodeQwenImage21",
            "inputs": {
                "clip": ["clip_loader", 0],
                "prompt": prompt,
                "negative_prompt": negative_prompt,
                "resolution": 0,
            },
        },
        "empty_latent": {"class_type": "EmptyLatentImage", "inputs": {"width": width, "height": height, "batch_size": 1}},
        "sampler": {
            "class_type": "KSampler",
            "inputs": {
                "model": ["unet_loader", 0],
                "positive": ["text_encode", 0],
                "negative": ["text_encode", 1],
                "latent_image": ["empty_latent", 0],
                "seed": seed,
                "steps": steps,
                "cfg": cfg,
                "sampler_name": "euler",
                "scheduler": "simple",
                "denoise": 1,
            },
        },
        "decode": {"class_type": "VAEDecode", "inputs": {"samples": ["sampler", 0], "vae": ["vae_loader", 0]}},
        "save": {"class_type": "SaveImage", "inputs": {"images": ["decode", 0], "filename_prefix": "gradio_qwen21"}},
    }


def _build_edit_graph(prompt: str, negative_prompt: str, image_filenames: list[str], steps: int, cfg: float, seed: int) -> dict:
    """Same pipeline as _build_t2i_graph, but TextEncodeQwenImage21 also
    receives up to 16 reference images (image_1..image_16, per ComfyUI's
    own nodes_qwen.py) and the empty latent's size comes from the VAE
    encoding the first reference image instead of a fixed width/height, so
    the edited output matches the first uploaded image's resolution.

    The "images" field is an Autogrow (dynamic) input, not a normal typed
    one - ComfyUI's API format does NOT accept it as a nested
    {"images": {"image_1": [...]}} dict under the node's own inputs. Its
    server-side parser (comfy_api/latest/_io.py, build_nested_inputs) only
    reconstructs that nesting from DOT-JOINED FLAT KEYS submitted directly
    in the node's inputs dict: "images.image_1", "images.image_2", etc.
    Submitting the nested shape (as an earlier version of this file did)
    is silently ignored - the node still runs, but with zero reference
    images attached, so the output becomes plain text-to-image with no
    connection to the uploaded photos. Confirmed against the real ComfyUI
    source, not assumed."""
    graph: dict = {
        "unet_loader": _unet_loader_node(),
        "clip_loader": {"class_type": "CLIPLoader", "inputs": {"clip_name": MODEL_FILES["clip"], "type": "qwen_image", "device": "default"}},
        "vae_loader": {"class_type": "VAELoader", "inputs": {"vae_name": MODEL_FILES["vae"]}},
    }

    # ComfyUI's own tokenizer (comfy/text_encoders/qwen_image21.py,
    # QwenImage21Tokenizer.tokenize_with_weights) already auto-prepends a
    # "<image1><vision tokens> <image2><vision tokens> {prompt}" template
    # whenever images are attached - this node does NOT need the caller to
    # add its own "Picture 1 / Picture 2" labels the way the older
    # TextEncodeQwenImageEditPlus node's execute() does in Python. Adding a
    # second, differently-worded numbering scheme on top would just add
    # redundant/conflicting signal ahead of the real <image1>/<image2>
    # anchor tokens. If the model still blends two reference people into
    # one, the fix is prompt wording (name which reference is which
    # explicitly, e.g. "the woman from image 1" / "the woman from image 2"),
    # not code - see the guide's Image Edit troubleshooting.
    text_encode_inputs = {
        "clip": ["clip_loader", 0],
        "vae": ["vae_loader", 0],
        "prompt": prompt,
        "negative_prompt": negative_prompt,
        "resolution": 1024,
    }
    for i, filename in enumerate(image_filenames[:16], start=1):
        load_id = f"load_image_{i}"
        graph[load_id] = {"class_type": "LoadImage", "inputs": {"image": filename}}
        text_encode_inputs[f"images.image_{i}"] = [load_id, 0]

    graph["text_encode"] = {
        "class_type": "TextEncodeQwenImage21",
        "inputs": text_encode_inputs,
    }
    graph["sampler"] = {
        "class_type": "KSampler",
        "inputs": {
            "model": ["unet_loader", 0],
            "positive": ["text_encode", 0],
            "negative": ["text_encode", 1],
            "latent_image": ["text_encode", 2],
            "seed": seed,
            "steps": steps,
            "cfg": cfg,
            "sampler_name": "euler",
            "scheduler": "simple",
            "denoise": 1,
        },
    }
    graph["decode"] = {"class_type": "VAEDecode", "inputs": {"samples": ["sampler", 0], "vae": ["vae_loader", 0]}}
    graph["save"] = {"class_type": "SaveImage", "inputs": {"images": ["decode", 0], "filename_prefix": "gradio_qwen21_edit"}}
    return graph


def make_generate_t2i(comfy_url_box: "gr.Textbox"):
    def generate_t2i(prompt: str, negative_prompt: str, width: int, height: int, steps: int, cfg: float, seed: int, randomize_seed: bool, comfy_url: str):
        if not prompt or not prompt.strip():
            raise gr.Error("Enter a prompt describing what to generate.")
        used_seed = random.randint(0, 2**32 - 1) if randomize_seed else int(seed)
        client = _client(comfy_url)
        graph = _build_t2i_graph(prompt.strip(), (negative_prompt or "").strip(), int(width), int(height), int(steps), float(cfg), used_seed)
        prompt_id = client.queue_prompt(graph)
        image = client.wait_for_result(prompt_id, save_node_id="save")
        return image, used_seed
    return generate_t2i


def generate_edit(prompt: str, reference_images: Optional[list], steps: int, cfg: float, seed: int, randomize_seed: bool, comfy_url: str):
    if not prompt or not prompt.strip():
        raise gr.Error("Enter a prompt describing the edit.")
    if not reference_images:
        raise gr.Error("Upload at least one reference image.")
    used_seed = random.randint(0, 2**32 - 1) if randomize_seed else int(seed)
    client = _client(comfy_url)
    filenames = []
    for item in reference_images:
        path = item[0] if isinstance(item, (list, tuple)) else item
        pil_image = Image.open(path).convert("RGB")
        filenames.append(client.upload_image(pil_image))
    graph = _build_edit_graph(prompt.strip(), "", filenames, int(steps), float(cfg), used_seed)
    prompt_id = client.queue_prompt(graph)
    image = client.wait_for_result(prompt_id, save_node_id="save")
    return image, used_seed


def generate_bg_removal(input_image, comfy_url: str):
    if input_image is None:
        raise gr.Error("Upload an image first.")
    client = _client(comfy_url)
    filename = client.upload_image(Image.open(input_image).convert("RGB") if isinstance(input_image, str) else input_image.convert("RGB"))
    seed = random.randint(0, 2**32 - 1)
    graph = _build_edit_graph(BG_REMOVAL_PROMPT, "", [filename], steps=25, cfg=1.0, seed=seed)
    prompt_id = client.queue_prompt(graph)
    image = client.wait_for_result(prompt_id, save_node_id="save")
    return image


def build_app(default_comfy_url: str) -> gr.Blocks:
    with gr.Blocks(title="Qwen-Image-2.1") as app:
        gr.Markdown(
            "# Qwen-Image-2.1\n"
            "Text to image, image editing, and background removal - no ComfyUI node graph. "
            "Runs against the ComfyUI instance below, on this same pod by default."
        )
        comfy_url = gr.Textbox(
            label="ComfyUI URL",
            value=default_comfy_url,
            info="Change this only if ComfyUI is running somewhere other than this same pod.",
        )

        with gr.Tab("Text to Image"):
            with gr.Row():
                with gr.Column():
                    t2i_prompt = gr.Textbox(label="Prompt", lines=3, placeholder="A red fox sitting in a snowy forest, soft morning light.")
                    t2i_negative = gr.Textbox(label="Negative prompt (optional)", lines=2, placeholder="blurry, low quality")
                    with gr.Row():
                        t2i_width = gr.Slider(label="Width", minimum=512, maximum=1536, step=32, value=1024)
                        t2i_height = gr.Slider(label="Height", minimum=512, maximum=1536, step=32, value=1024)
                    with gr.Row():
                        t2i_steps = gr.Slider(label="Steps", minimum=10, maximum=50, step=1, value=25)
                        t2i_cfg = gr.Slider(label="CFG", minimum=1.0, maximum=2.0, step=0.05, value=1.0)
                    with gr.Row():
                        t2i_seed = gr.Number(label="Seed", value=0, precision=0)
                        t2i_randomize = gr.Checkbox(label="Randomize seed", value=True)
                    t2i_button = gr.Button("Generate", variant="primary")
                with gr.Column():
                    t2i_output = gr.Image(label="Result", type="pil")
                    t2i_used_seed = gr.Number(label="Seed used", precision=0, interactive=False)
            t2i_button.click(
                fn=make_generate_t2i(comfy_url),
                inputs=[t2i_prompt, t2i_negative, t2i_width, t2i_height, t2i_steps, t2i_cfg, t2i_seed, t2i_randomize, comfy_url],
                outputs=[t2i_output, t2i_used_seed],
            )

        with gr.Tab("Image Edit"):
            with gr.Row():
                with gr.Column():
                    edit_upload = gr.File(
                        label="Reference images (up to 16) - select several at once, or add more later",
                        file_count="multiple",
                        file_types=["image"],
                        type="filepath",
                    )
                    edit_images_preview = gr.Gallery(label="Selected", type="filepath", columns=4, height=200, preview=False)
                    edit_clear_button = gr.Button("Clear all reference images", size="sm")
                    edit_prompt = gr.Textbox(
                        label="Prompt",
                        lines=3,
                        placeholder="Keep the person from image 1 unchanged, put the jacket from image 2 on them, set the scene in an autumn park.",
                    )
                    with gr.Row():
                        edit_steps = gr.Slider(label="Steps", minimum=10, maximum=50, step=1, value=25)
                        edit_cfg = gr.Slider(label="CFG", minimum=1.0, maximum=2.0, step=0.05, value=1.0)
                    with gr.Row():
                        edit_seed = gr.Number(label="Seed", value=0, precision=0)
                        edit_randomize = gr.Checkbox(label="Randomize seed", value=True)
                    edit_button = gr.Button("Generate", variant="primary")
                with gr.Column():
                    edit_output = gr.Image(label="Result", type="pil")
                    edit_used_seed = gr.Number(label="Seed used", precision=0, interactive=False)

            def _accumulate_images(new_files, existing):
                """gr.File's own value already replaces on each selection (a
                browser file picker has no memory of a prior selection), so
                without this the second upload would silently drop the
                first. Keep every filepath seen so far instead, deduplicated
                so re-selecting the same file twice doesn't double it up."""
                existing = existing or []
                existing_paths = [p for p, _ in existing] if existing and isinstance(existing[0], (list, tuple)) else list(existing)
                combined = list(existing_paths)
                for f in (new_files or []):
                    path = f if isinstance(f, str) else getattr(f, "name", None)
                    if path and path not in combined:
                        combined.append(path)
                return combined[:16], combined[:16]

            edit_upload.change(fn=_accumulate_images, inputs=[edit_upload, edit_images_preview], outputs=[edit_images_preview, edit_images_preview])
            edit_clear_button.click(fn=lambda: (None, []), inputs=None, outputs=[edit_upload, edit_images_preview])
            edit_button.click(
                fn=generate_edit,
                inputs=[edit_prompt, edit_images_preview, edit_steps, edit_cfg, edit_seed, edit_randomize, comfy_url],
                outputs=[edit_output, edit_used_seed],
            )

        with gr.Tab("Remove Background"):
            with gr.Row():
                with gr.Column():
                    bg_input = gr.Image(label="Image", type="pil")
                    bg_button = gr.Button("Remove background", variant="primary")
                with gr.Column():
                    bg_output = gr.Image(label="Result (transparent PNG)", type="pil", image_mode="RGBA")
            bg_button.click(fn=generate_bg_removal, inputs=[bg_input, comfy_url], outputs=bg_output)

        gr.Markdown(
            "Qwen-Image-2.1's weights are licensed for non-commercial (research/evaluation) use only. "
            "See the full setup guide for the license details and troubleshooting."
        )
    return app


def main() -> None:
    parser = argparse.ArgumentParser(description="Qwen-Image-2.1 Gradio UI")
    parser.add_argument("--comfy-url", default=os.environ.get("COMFY_URL", "http://127.0.0.1:8188"), help="Base URL of the ComfyUI instance to drive")
    parser.add_argument("--host", default="0.0.0.0")
    parser.add_argument("--port", type=int, default=8889)
    parser.add_argument("--unet", default=None, help="Override the diffusion model filename (must already be in ComfyUI/models/diffusion_models/)")
    parser.add_argument("--clip", default=None, help="Override the text encoder filename")
    parser.add_argument("--vae", default=None, help="Override the VAE filename")
    args = parser.parse_args()

    if args.unet:
        MODEL_FILES["unet"] = args.unet
    if args.clip:
        MODEL_FILES["clip"] = args.clip
    if args.vae:
        MODEL_FILES["vae"] = args.vae

    app = build_app(args.comfy_url)
    app.queue(max_size=10).launch(server_name=args.host, server_port=args.port)


if __name__ == "__main__":
    main()