app.py

"""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.

The diffusion pipeline stays resident on `cuda` (loaded at module level, as
ZeroGPU requires). The 9B prompt-rewriting models don't fit alongside it on a
48GB card, so they live in CPU RAM and are moved onto the GPU only inside
their @spaces.GPU calls. 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

# 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")

# --- 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).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).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 = 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)


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


@spaces.GPU(duration=300)
def enhance_prompt_t2i(prompt: str) -> str:
    """Rewrite a brief text-to-image request into a detailed English prompt
    with Qwen-Image-2.1-PE-T2I."""
    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,
    )
    pe_t2i.to("cuda")
    try:
        inputs = pe_t2i_tokenizer(text, return_tensors="pt").to("cuda")
        with torch.no_grad():
            out = pe_t2i.generate(
                **inputs, max_new_tokens=4096,
                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
        )
    finally:
        pe_t2i.to("cpu")
        torch.cuda.empty_cache()
    return _parse_rewrite(gen, prompt)


@spaces.GPU(duration=300)
def enhance_prompt_i2i(image, instruction: str) -> str:
    """Rewrite an image-editing instruction against the condition image with
    Qwen-Image-2.1-PE-I2I."""
    _check_prompt_guard(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},
        ]},
    ]
    pe_i2i.to("cuda")
    try:
        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=4096,
                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
        )
    finally:
        pe_i2i.to("cpu")
        torch.cuda.empty_cache()
    return _parse_rewrite(gen, instruction)


@spaces.GPU(duration=180)
def text_to_image(prompt: str, steps: int = 40) -> 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(duration=180)
def edit_image(image, instruction: str, steps: int = 40) -> 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()