app.py

"""Qwen-Image-2.1 — unified text-to-image generation and image editing on ZeroGPU.

Pipeline for one request:

1. NCII guard (CPU, sidecar subprocess) — image-input requests have their prompt screened
   with ``hfmlsoc/ncii-guard-v02`` before any GPU time is reserved.
2. Optional prompt enhancement — delegated over the Gradio API to a companion Space
   that keeps ``Qwen-Image-2.1-PE-T2I`` and ``Qwen-Image-2.1-PE-I2I`` warm.
3. Generation — ``QwenImage21Pipeline`` inside ``@spaces.GPU``.

Steps 1 and 2 deliberately run outside the GPU fork so a rejected or enhanced prompt
costs the visitor no ZeroGPU quota.
"""

import os

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import spaces  # noqa: E402  — must precede torch / any CUDA-touching import

import json  # noqa: E402
import math  # noqa: E402
import random  # noqa: E402
import tempfile  # noqa: E402
import time  # noqa: E402

import gradio as gr  # noqa: E402
import torch  # noqa: E402
from diffusers import QwenImage21Pipeline  # noqa: E402
from PIL import Image  # noqa: E402
import ncii_guard  # noqa: E402
from qwen21_aoti import aoti_load_pipeline  # noqa: E402

MODEL_ID = "Qwen/Qwen-Image-2.1"
AOTI_REPO = os.environ.get("QWEN21_AOTI_REPO", "hugging-apps/qwen-image-2-1-aoti")
GUARD_ID = ncii_guard.GUARD_ID
PE_SPACE_ID = os.environ.get("PE_SPACE_ID", "hugging-apps/qwen-image-2-1-prompt-enhancer")
# The enhancer budgets its ZeroGPU reservation from this cap (~23 tok/s), so the cap is what a
# visitor's quota is charged, not the time actually spent. Measured on the enhancer Space: the
# I2I rewriter runs to the cap whether it is 1024 (≈39s, reservation 56s) or 2048 (≈78s,
# reservation 108s), so doubling it does not stop it being cut off — it only doubles the charge.
# A cut-off answer is salvaged into a usable prompt (minus the ratio hint), and on the rare total
# failure ``prepare`` warns and falls back to the user's own prompt. So keep the cap lean.
PE_MAX_NEW_TOKENS = int(os.environ.get("PE_MAX_NEW_TOKENS", "1024"))

MAX_SEED = 2**31 - 1
MAX_INPUT_IMAGES = 10
EXAMPLE_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "examples")

# Shown verbatim when the guard rejects a prompt — deliberately non-explicit.
GUARD_REJECTION_MESSAGE = "prompt invalid based on our classifiers, try again"
# See the model card's threshold sweep: 0.5 is the balanced precision/recall point.
GUARD_THRESHOLD = 0.5

# The aspect ratios the model card lists, at its 2048-base resolution.
BASE_RESOLUTION = 2048
ASPECT_RATIOS = {
    "1:1": (2048, 2048),
    "4:3": (2400, 1792),
    "3:4": (1792, 2400),
    "3:2": (2528, 1696),
    "2:3": (1696, 2528),
    "16:9": (2752, 1536),
    "9:16": (1536, 2752),
}
RESOLUTION_CHOICES = [("1K — fastest", 1024), ("1.5K", 1536), ("2K — model default", 2048)]
# The step-cost fit below was measured with the AOTI kernels active; falling back to eager
# (QWEN21_AOTI=0, or a failed artifact load) costs roughly a third more per step.
EAGER_STEP_PENALTY = 1.3
KV_CACHE_BYTES_PER_TOKEN = 32 * 2 * 4096 * 2
KV_CACHE_BUDGET_GB = 10.0


def prefix_tokens(n_images: int, resolution: int) -> int:
    return n_images * (int(resolution) // 16) ** 2


def kv_cache_fits(n_images: int, resolution: int) -> bool:
    return prefix_tokens(n_images, resolution) * KV_CACHE_BYTES_PER_TOKEN <= KV_CACHE_BUDGET_GB * 1e9

TRANSPARENCY_INFO = (
    "Describe what to generate, or how to edit the input images (refer to them as "
    "<image1>…<image10> in upload order). For a transparent RGBA result, phrase the prompt as: "
    "“This is an RGBA image with transparency. … The image has alpha channel and the "
    "background is transparent.”"
)

# --------------------------------------------------------------------------------------
# Models
# --------------------------------------------------------------------------------------

print(f"[load] {MODEL_ID} …", flush=True)
pipe = QwenImage21Pipeline.from_pretrained(MODEL_ID, torch_dtype=torch.bfloat16)
pipe = pipe.to("cuda")

# The VAE decoder upcasts to fp32 in its upsamplers, so a 2048x2048 decode allocates far
# more than the denoising loop that produced it: it died with an allocator assert
# ("NVML_SUCCESS == r INTERNAL ASSERT FAILED") *after* all 40 steps had completed, and
# left the ZeroGPU worker wedged so every later request reported "No CUDA GPUs are
# available" until the Space was restarted. Tiling the decode keeps peak VAE memory flat.
# The threshold is set above 1536 so the resolutions measured to decode safely in one
# shot still do, and only larger outputs pay the tiling overhead.
pipe.vae.enable_tiling(
    tile_sample_min_height=1536,
    tile_sample_min_width=1536,
    tile_sample_stride_height=1152,
    tile_sample_stride_width=1152,
)

AOTI_LOADED = []
if os.environ.get("QWEN21_AOTI", "1") != "0":
    try:
        AOTI_LOADED, _aoti_config = aoti_load_pipeline(pipe, AOTI_REPO, token=os.environ.get("HF_TOKEN"))
        print(f"[load] aoti kernels from {AOTI_REPO}: {AOTI_LOADED} (torch {_aoti_config['torch']})", flush=True)
    except Exception as exc:  # noqa: BLE001
        print(f"[load] aoti unavailable, running eager: {exc!r}", flush=True)
AOTI_BLOCKS = "QwenImage21DecodeBlock" in AOTI_LOADED
print("[load] pipeline ready", flush=True)

print(f"[load] {GUARD_ID} sidecar …", flush=True)
ncii_guard.start()
print("[load] guard ready", flush=True)


def guard_score(text: str) -> float:
    """Probability that ``text`` is a request for non-consensual intimate imagery."""
    return ncii_guard.score(text or "")


def check_guard(prompt: str) -> None:
    """Raise the non-explicit rejection error when the guard flags ``prompt``."""
    score = guard_score(prompt)
    print(f"[guard] score={score:.3f} threshold={GUARD_THRESHOLD}", flush=True)
    if score >= GUARD_THRESHOLD:
        raise gr.Error(GUARD_REJECTION_MESSAGE)


# --------------------------------------------------------------------------------------
# Prompt enhancement (delegated to the companion Space)
# --------------------------------------------------------------------------------------

_pe_client = None


def get_pe_client():
    """Lazily connect to the prompt-enhancement Space (its models are already warm)."""
    global _pe_client
    if _pe_client is None:
        from gradio_client import Client

        token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")
        # gradio_client 2.x renamed ``hf_token`` to ``token``; accept either client version.
        try:
            _pe_client = Client(
                PE_SPACE_ID, token=token, httpx_kwargs={"timeout": 900}, verbose=False
            )
        except TypeError:
            _pe_client = Client(
                PE_SPACE_ID, hf_token=token, httpx_kwargs={"timeout": 900}, verbose=False
            )
    return _pe_client


def enhance_prompt(prompt: str, image_paths: list[str]) -> tuple[str, str]:
    """Rewrite ``prompt`` with the official PE models. Returns (prompt, wh_ratio)."""
    from gradio_client import handle_file

    client = get_pe_client()
    started = time.perf_counter()
    result = client.predict(
        prompt=prompt,
        image_paths=[handle_file(p) for p in image_paths],
        max_new_tokens=PE_MAX_NEW_TOKENS,
        enable_thinking=False,
        seed=0,
        randomize_seed=True,
        api_name="/enhance",
    )
    rewritten, wh_ratio = str(result[0]).strip(), str(result[1]).strip()
    print(f"[enhance] {time.perf_counter() - started:.1f}s -> ratio={wh_ratio!r}", flush=True)
    return (rewritten or prompt), wh_ratio


# --------------------------------------------------------------------------------------
# Helpers
# --------------------------------------------------------------------------------------


def normalize_gallery(gallery) -> list[str]:
    """Reduce a ``gr.Gallery`` value to a flat list of local file paths."""
    paths: list[str] = []
    for item in gallery or []:
        if isinstance(item, (tuple, list)):
            item = item[0]
        if isinstance(item, dict):
            item = item.get("path") or item.get("image") or item.get("name")
            if isinstance(item, dict):
                item = item.get("path")
        if isinstance(item, Image.Image):
            tmp = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
            item.save(tmp.name)
            item = tmp.name
        if isinstance(item, str) and item:
            paths.append(item)
    return paths


def resolve_size(resolution: int, aspect_ratio: str) -> tuple[int | None, int | None]:
    """Map a resolution tier + aspect ratio onto the model card's width/height pairs."""
    if not aspect_ratio or aspect_ratio == "Auto" or aspect_ratio not in ASPECT_RATIOS:
        return None, None
    base_w, base_h = ASPECT_RATIOS[aspect_ratio]
    scale = int(resolution) / BASE_RESOLUTION
    round32 = lambda v: max(32, int(round(v * scale / 32)) * 32)  # noqa: E731
    return round32(base_w), round32(base_h)


def estimate_duration(
    prompt="",
    input_images=None,
    negative_prompt="",
    true_cfg_scale=1.0,
    num_inference_steps=40,
    seed=0,
    resolution=1024,
    aspect_ratio="Auto",
    *args,
    **kwargs,
):
    """Budget GPU seconds from times measured on this Space's own hardware.

    Refitted against ``[generate]`` log timings with the AOTI kernels live, all at 40 steps:
    1024px = 16.3s (18.3s worst), 1536px = 40.1s, 2048px = 82.9s, and 1024px with one and two
    condition images = 21.7s and 23.1s. Step cost scales as ``latent_tokens ** 1.363``; the
    cached prefix of a condition image contributes about a fifth of its tokens to each step on
    top of a couple of seconds of one-off prefill.

    True CFG doubles the step (it denoises the negative branch too) but the pipeline only turns
    it on when a negative prompt is actually set — see ``_diffuse`` — so this mirrors that
    condition rather than keying off the scale alone. Doing otherwise doubled the reservation of
    every default request, which is charged to the visitor's ZeroGPU quota whether or not it is
    spent. 1.25 is the only slack, which lands ~10-30% above the measured call.
    """
    try:
        steps = int(num_inference_steps)
        res = int(resolution)
        cfg = float(true_cfg_scale)
    except (TypeError, ValueError):
        steps, res, cfg = 40, 1024, 1.0
    n_images = len(input_images or [])
    tokens = (res / 16.0) ** 2
    prefix = prefix_tokens(n_images, res)
    cached = kv_cache_fits(n_images, res)
    tokens += 0.2 * prefix if cached else prefix
    per_step = 3.4e-6 * tokens**1.363
    if cfg > 1.0 and (negative_prompt or "").strip():
        per_step *= 2.0
    if not (AOTI_BLOCKS and cached):
        per_step *= EAGER_STEP_PENALTY
    fixed = 4.0 + 2.0 * n_images + 1.5e-6 * res * res
    return int(min(340, math.ceil((fixed + steps * per_step) * 1.25)))


# --------------------------------------------------------------------------------------
# Inference
# --------------------------------------------------------------------------------------


@spaces.GPU(duration=estimate_duration)
def _diffuse(
    prompt: str,
    image_paths: list,
    negative_prompt: str,
    true_cfg_scale: float,
    num_inference_steps: int,
    seed: int,
    resolution: int,
    aspect_ratio: str,
    progress=gr.Progress(track_tqdm=True),
):
    """Run the diffusion pipeline. Inputs are already validated and screened."""
    condition_images = [Image.open(p) for p in image_paths] or None
    width, height = resolve_size(resolution, aspect_ratio)

    negative_prompt = (negative_prompt or "").strip()
    use_kv_cache = kv_cache_fits(len(image_paths), int(resolution))
    call_kwargs = {"use_kv_cache": use_kv_cache}
    if negative_prompt and float(true_cfg_scale) > 1.0:
        call_kwargs["negative_prompt"] = negative_prompt
        call_kwargs["true_cfg_scale"] = float(true_cfg_scale)

    started = time.perf_counter()
    image = pipe(
        prompt=prompt,
        image=condition_images,
        height=height,
        width=width,
        output_resolution=int(resolution),
        num_inference_steps=int(num_inference_steps),
        generator=torch.Generator(device="cuda").manual_seed(int(seed) % (MAX_SEED + 1)),
        **call_kwargs,
    ).images[0]
    print(
        f"[generate] images={len(image_paths)} steps={num_inference_steps} "
        f"res={resolution} ratio={aspect_ratio} size={image.size} kv_cache={use_kv_cache} "
        f"aoti={AOTI_LOADED} "
        f"elapsed={time.perf_counter() - started:.1f}s",
        flush=True,
    )
    return image


def generate_image(
    prompt: str = "",
    input_images=None,
    negative_prompt: str = "",
    true_cfg_scale: float = 1.0,
    num_inference_steps: int = 40,
    seed: int = 0,
    resolution: int = 1024,
    aspect_ratio: str = "Auto",
    progress=gr.Progress(track_tqdm=True),
):
    """Generate or edit an image with Qwen-Image-2.1.

    With no input images this is text-to-image; with one to ten input images it is
    image editing, and the prompt may refer to them as <image1>…<image10>.

    Args:
        prompt: What to generate, or how to edit the inputs.
        input_images: Up to 10 condition images. Leave empty for text-to-image.
        negative_prompt: What to avoid. Only applied when true_cfg_scale > 1.
        true_cfg_scale: Classifier-free guidance scale; 1.0 disables it.
        num_inference_steps: Number of flow-matching steps.
        seed: Random seed.
        resolution: Target side length in pixels (1024, 1536 or 2048).
        aspect_ratio: One of the model's supported ratios, or "Auto".

    Returns:
        The generated image, as RGBA (the model natively supports transparency).
    """
    # Validation and screening deliberately happen here, in the main process: a rejected
    # prompt must never reserve a GPU slot, and gr.Error raised inside the ZeroGPU worker
    # reaches the browser as a bare "RuntimeError" instead of its message.
    prompt = (prompt or "").strip()
    if not prompt:
        raise gr.Error("Please enter a prompt.")

    image_paths = normalize_gallery(input_images)
    if len(image_paths) > MAX_INPUT_IMAGES:
        raise gr.Error(f"Up to {MAX_INPUT_IMAGES} input images are supported.")

    # `prepare` already screened the UI path; this keeps the guard in force for callers
    # that hit /generate directly.
    if image_paths:
        check_guard(prompt)

    # Gradio only binds a live gr.Progress to the function it calls, which is this one --
    # the copy sitting in _diffuse's signature is an inert import-time default, so without
    # forwarding this the UI showed no step counter at all during a 25-120s generation.
    return _diffuse(
        prompt,
        image_paths,
        negative_prompt,
        true_cfg_scale,
        num_inference_steps,
        seed,
        resolution,
        aspect_ratio,
        progress,
    )


def prepare(prompt, input_images, enhance, seed, randomize_seed, aspect_ratio):
    """Screen the prompt, optionally rewrite it, and resolve the seed — all off-GPU."""
    prompt = (prompt or "").strip()
    if not prompt:
        raise gr.Error("Please enter a prompt.")

    image_paths = normalize_gallery(input_images)
    if len(image_paths) > MAX_INPUT_IMAGES:
        raise gr.Error(f"Up to {MAX_INPUT_IMAGES} input images are supported.")

    # Guard image-input requests before any GPU time is reserved.
    if image_paths:
        score = guard_score(prompt)
        print(f"[guard] score={score:.3f} threshold={GUARD_THRESHOLD}", flush=True)
        if score >= GUARD_THRESHOLD:
            raise gr.Error(GUARD_REJECTION_MESSAGE)

    resolved_seed = random.randint(0, MAX_SEED) if randomize_seed else int(seed)

    final_prompt, rewritten_display = prompt, ""
    if enhance:
        try:
            final_prompt, wh_ratio = enhance_prompt(prompt, image_paths)
            rewritten_display = final_prompt
            # Honour the rewriter's ratio only when the user left the choice open and
            # there is no input image whose aspect ratio the pipeline should follow.
            if wh_ratio and aspect_ratio == "Auto" and not image_paths:
                if wh_ratio in ASPECT_RATIOS:
                    aspect_ratio = wh_ratio
        except Exception as exc:  # noqa: BLE001 — never let PE take the whole request down
            print(f"[enhance] failed: {exc!r}", flush=True)
            gr.Warning(f"Prompt enhancement unavailable, using the original prompt ({exc}).")

    return final_prompt, rewritten_display, resolved_seed, aspect_ratio


# --------------------------------------------------------------------------------------
# Examples (Qwen's own demo cases, bundled under examples/)
# --------------------------------------------------------------------------------------

with open(os.path.join(EXAMPLE_DIR, "cases.json"), encoding="utf-8") as fh:
    DEMO_CASES = json.load(fh)

T2I_EXAMPLES = [[c["prompt"]] for c in DEMO_CASES if not c["inputs"]]
T2I_LABELS = [c["title_en"] for c in DEMO_CASES if not c["inputs"]]
EDIT_EXAMPLES = [
    [c["prompt"], [os.path.join(EXAMPLE_DIR, name) for name in c["inputs"]]]
    for c in DEMO_CASES
    if c["inputs"]
]
EDIT_LABELS = [c["title_en"] for c in DEMO_CASES if c["inputs"]]


# --------------------------------------------------------------------------------------
# UI
# --------------------------------------------------------------------------------------

css = """
#col-container { margin: 0 auto; max-width: 1100px; }
#logo, #logo > * { padding: 0 !important; margin: 0 !important; min-height: 0 !important; border: none !important; background: transparent !important; }
#logo img { display: block; margin: 0 auto !important; }
/* Citrus leaves some text colours unset in dark mode, which renders dark-on-dark. */
.dark .gradio-container { color: var(--body-text-color); }
"""

# `theme` / `css` are passed to launch(), not to the Blocks constructor: Gradio 6 moved them
# and warns ("The parameters have been moved from the Blocks constructor to the launch()
# method in Gradio 6.0: theme, css") if they are set here.
with gr.Blocks(title="Qwen-Image-2.1") as demo:
    with gr.Column(elem_id="col-container"):
        gr.HTML(
            '<img src="https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-Image/image2.1/logo.png"'
            ' alt="Qwen-Image logo" width="380" style="display:block;margin:0 auto;">',
            elem_id="logo",
        )

        with gr.Row():
            with gr.Column():
                input_images = gr.Gallery(
                    label="Input images (editing, up to 10)",
                    type="filepath",
                    interactive=True,
                    columns=5,
                    height=300,
                    show_label=True,
                )
            with gr.Column():
                result = gr.Image(
                    label="Result",
                    type="pil",
                    image_mode="RGBA",
                    format="png",
                    height=300,
                    show_label=True,
                )

        prompt = gr.Textbox(
            label="Prompt",
            info=TRANSPARENCY_INFO,
            lines=5,
            max_lines=24,
            placeholder="Describe what you want to generate or edit…",
        )

        with gr.Row():
            enhance = gr.Checkbox(label="Enhance prompt", value=True, scale=1)
            generate_button = gr.Button("Generate image", variant="primary", scale=2)

        rewritten_prompt_output = gr.Textbox(
            label="Rewritten prompt",
            lines=4,
            max_lines=20,
            interactive=False,
            placeholder="With prompt enhancement on, the rewritten prompt appears here.",
        )

        with gr.Accordion("Advanced settings", open=False):
            with gr.Row():
                resolution = gr.Dropdown(
                    label="Resolution", choices=RESOLUTION_CHOICES, value=1024
                )
                aspect_ratio = gr.Dropdown(
                    label="Aspect ratio",
                    choices=["Auto"] + list(ASPECT_RATIOS),
                    value="Auto",
                    info="Auto follows the input image, or squares up for text-to-image.",
                )
            num_inference_steps = gr.Slider(
                label="Inference steps", minimum=8, maximum=60, step=1, value=40
            )
            with gr.Row():
                seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0)
                randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
            negative_prompt = gr.Textbox(
                label="Negative prompt",
                value="",
                lines=2,
                info="Only applied when guidance is above 1.0.",
            )
            true_cfg_scale = gr.Slider(
                label="True CFG scale", minimum=1.0, maximum=10.0, step=0.1, value=1.0
            )

        # A hidden Textbox rather than gr.State: it carries the enhanced prompt from
        # `prepare` into `generate_image`, and being a real component it is also what
        # makes `prompt` a first-class parameter of the /generate API endpoint.
        final_prompt_box = gr.Textbox(visible=False, label="Prompt")

        gr.Markdown("### Examples")
        gr.Markdown("**Text to image**")
        gr.Examples(
            examples=T2I_EXAMPLES,
            example_labels=T2I_LABELS,
            inputs=[prompt],
            fn=generate_image,
            outputs=[result],
            cache_examples=True,
            cache_mode="lazy",
            label="Text-to-image cases",
        )
        gr.Markdown("**Image editing** — the example loads its reference images too")
        gr.Examples(
            examples=EDIT_EXAMPLES,
            example_labels=EDIT_LABELS,
            inputs=[prompt, input_images],
            fn=generate_image,
            outputs=[result],
            cache_examples=True,
            cache_mode="lazy",
            label="Image-editing cases",
        )

    generate_button.click(
        fn=prepare,
        inputs=[prompt, input_images, enhance, seed, randomize_seed, aspect_ratio],
        outputs=[final_prompt_box, rewritten_prompt_output, seed, aspect_ratio],
    ).then(
        fn=generate_image,
        inputs=[
            final_prompt_box,
            input_images,
            negative_prompt,
            true_cfg_scale,
            num_inference_steps,
            seed,
            resolution,
            aspect_ratio,
        ],
        outputs=[result],
        api_name="generate",
    )


demo.queue(default_concurrency_limit=1, max_size=20)

if __name__ == "__main__":
    demo.launch(
        mcp_server=True,
        theme=gr.themes.Citrus(),
        css=css,
        allowed_paths=[EXAMPLE_DIR],
    )