app.py

"""Gradio demo for the DMD-distilled Qwen-Image-2.1 student (v0.2.1: LoRA r256 sampled in 6 steps; v0.2 / v0.1 students still selectable)."""

# ZeroGPU patches torch at import time, so `spaces` must be imported before torch or any CUDA use.
# find_spec keeps this file runnable off-Spaces, where the package is absent.
import importlib.util
import os

# Less allocator fragmentation (untiled 2048^2 text-to-image reserves ~65 GiB of the 96 GB card).
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

if importlib.util.find_spec("spaces"):
    import spaces

    # size="xlarge" (the full 96 GB RTX Pro 6000 Blackwell, 2x quota): the pipeline runs exactly as validated, with no VAE
    # tiling anywhere - tiling the encode wrecks reference-conditioned edits (duplicated subjects, wrong scale), and on the
    # 48 GB "large" card (44.7 GiB usable) the untiled calls do not fit: three-reference edit at a 1024^2-area output
    # 44.3 GiB, 2048^2 text-to-image 57.0 GiB, three-reference edit at 1760x1344 51.5 GiB (measured, B200, bf16 - see README).
    # duration=90 covers the prompt rewrite plus the largest call.
    gpu = spaces.GPU(duration=90, size="xlarge")
else:
    gpu = lambda fn: fn  # noqa: E731

import json
import random
import re
import time
from pathlib import Path

import gradio as gr
import torch
import compare
from PIL import Image
from diffusers import FlowMatchEulerDiscreteScheduler, QwenImage21Pipeline, QwenImage21Transformer2DModel
from diffusers.pipelines.qwenimage21.pipeline_qwenimage21 import calculate_dimensions

# Students in the model repo. STUDENT="lora" (default): the v0.2.1 LoRA (r256) at the repo root, applied at runtime on
# top of the base transformer and never merged (merging into bf16 keeps only ~47% of the adapter delta, the
# per-element deltas sit below the bf16 ULP of the base weights). LORA_FILE selects the adapter file; the v0.2 and
# v0.1 adapters are still in the repo. STUDENT="full": the v0.1 full fine-tune in transformer/ (bf16, exact, no
# adapter) replaces the base transformer at load time - the base transformer is then not downloaded at all.
BASE_MODEL_ID = os.environ.get("BASE_MODEL_ID", "Qwen/Qwen-Image-2.1")
STUDENT_REPO = os.environ.get("STUDENT_REPO", os.environ.get("LORA_REPO", "Viggle/Qwen-Image-2.1-viggle-turbo"))
STUDENT = os.environ.get("STUDENT", "lora")
LORA_FILE = os.environ.get("LORA_FILE", "Qwen-Image-2.1-viggle-turbo-v0.2.1-6step-lora-r256.safetensors")
STUDENT_TAG = {"full": "v0.1 full fine-tune (transformer/)", "lora": "v0.2.1 LoRA r256" if "v0.2.1" in LORA_FILE else f"LoRA {LORA_FILE}"}[STUDENT]
HF_TOKEN = os.environ.get("HF_TOKEN")  # only needed while the model repo is private
# The LoRA is sampled on 6 Euler steps whose raw (pre-shift) sigma nodes are RAW_NODES: the 4-step training schedule
# linspace(1, 1/4, 4) with its highest-noise segment [1, 0.75] cut into three (1, 0.9375, 0.875, 0.75). The composition
# is decided in that segment, and one big Euler step there ghosts and drifts the layout; the low-noise nodes 0.75, 0.5,
# 0.25 are the ones the student was trained to land on, and moving them softens the image. So every extra step goes to the high-noise
# end: raw_nodes(steps) subdivides [1, 0.875] evenly and never touches 0.875, 0.75, 0.5, 0.25 (5 = the schedule v0.2
# launched with, 6 = default, 7 also validated; 4 = the training nodes; 3 is the pipeline's default linspace and
# off-contract). Passing sigmas= to the pipeline is the whole recipe - it applies the same resolution-dependent
# shift it applies to its default nodes.
STEPS = 6


def raw_nodes(steps):
    """Raw sigma nodes for a step count: the 4-step training nodes, all extra steps in the highest-noise segment."""
    if steps < 4:
        return None  # pipeline default linspace(1, 1/steps, steps); off-contract
    if steps == 4:
        return [1.0, 0.75, 0.5, 0.25]
    return [1.0 - 0.125 * i / (steps - 4) for i in range(steps - 4)] + [0.875, 0.75, 0.5, 0.25]


RAW_NODES = raw_nodes(STEPS)
MAX_SEED = 2**31 - 1
EXAMPLES_DIR = Path(__file__).parent / "examples"
# examples/manifest.json is written by release/render_examples.py, which runs *this* app.py's generate()
# over every row (STEPS steps, prompt enhancement on, fixed seed) and stores the result next to the references,
# so the examples table shows real outputs of this model without spending GPU time on a click.
EXAMPLES = json.loads((EXAMPLES_DIR / "manifest.json").read_text(encoding="utf-8")) if (EXAMPLES_DIR / "manifest.json").exists() else []

# Output sizes are calculate_dimensions(area, ratio) — the pipeline's own rule, rounded to multiples
# of 32 — over the target areas the student was distilled on. build_user_manifest.py buckets targets
# at 1024^2 / 2048^2 for text-to-image and 1024^2 / 1536^2 for editing, so the two modes get
# different menus. Reference images are always encoded at 1024^2 area regardless.
RATIOS = [("1:1", 1.0), ("16:9", 16 / 9), ("9:16", 9 / 16), ("4:3", 4 / 3), ("3:4", 3 / 4), ("3:2", 3 / 2), ("2:3", 2 / 3)]
AUTO = "Auto · match the last reference (1024² area)"


def _bucket(area):
    sizes = {}
    for name, ratio in RATIOS:
        width, height, _ = calculate_dimensions(area * area, ratio)
        sizes[f"{name} · {width}×{height} ({area}² area)"] = (width, height)
    return sizes


SIZES = {AUTO: None, **_bucket(1024), **_bucket(1536), **_bucket(2048)}
T2I_CHOICES = [*_bucket(1024), *_bucket(2048)]
EDIT_CHOICES = [AUTO, *_bucket(1024), *_bucket(1536)]
# the (width, height) pairs the editing menu actually offers, in menu order
EDIT_DIMS = [SIZES[label] for label in EDIT_CHOICES if SIZES[label]]

if STUDENT == "full":
    transformer = QwenImage21Transformer2DModel.from_pretrained(
        STUDENT_REPO, subfolder="transformer", torch_dtype=torch.bfloat16, token=HF_TOKEN
    )
    pipe = QwenImage21Pipeline.from_pretrained(BASE_MODEL_ID, transformer=transformer, dtype=torch.bfloat16, token=HF_TOKEN)
else:
    pipe = QwenImage21Pipeline.from_pretrained(BASE_MODEL_ID, dtype=torch.bfloat16, token=HF_TOKEN)
    pipe.load_lora_weights(STUDENT_REPO, weight_name=LORA_FILE, token=HF_TOKEN)
# The stock Qwen-Image-2.1 scheduler config carries shift_terminal=0.02, which stretches the last
# sigma node to 0.02 instead of 0 and silently costs the final step. The student was distilled
# against the unstretched schedule, so rebuild the scheduler with shift_terminal=None; then
# pipe(num_inference_steps=4) alone reproduces dmd_common.student_sigmas(terminal="none") exactly.
pipe.scheduler = FlowMatchEulerDiscreteScheduler.from_config(pipe.scheduler.config, shift_terminal=None)
pipe.to("cuda")


# transformers' Qwen3VLVisionPatchEmbed runs an nn.Conv3d whose kernel_size equals its stride, which is
# exactly a linear map over each flattened patch. cuDNN has no usable bf16 Conv3d kernel for this shape and
# falls back to one that costs ~30 s per reference image (measured; 356 ms even with memory to spare, versus
# 2.9 ms for the same conv in fp32). The matmul below is 0.8 ms and its output is bitwise equal to the conv's.
patch_embed = pipe.text_encoder.model.visual.patch_embed


def _patch_embed_forward(hidden_states):
    proj = patch_embed.proj
    weight = proj.weight.reshape(proj.weight.shape[0], -1)
    return hidden_states.to(weight.dtype).flatten(1) @ weight.T + proj.bias


patch_embed.forward = _patch_embed_forward

# Prompt enhancement reuses the pipeline's own text encoder: Qwen-Image-2.1 ships the complete
# Qwen3-VL-8B-Instruct (LM head included) as its text encoder, so it can rewrite prompts with the system
# prompts published with Qwen/Qwen-Image-2.1-PE-T2I and -PE-I2I. The official rewriters themselves
# (Qwen3.5-VL 9B fine-tunes, ~18 GiB) would not fit next to the pipeline on a 48 GB card, and the
# official Space calls a hosted API instead. This costs no extra weights, ~2 GiB of KV cache and
# 4-11 s per call (B200); greedy decoding, so a given prompt always rewrites the same way.
PE_T2I = (Path(__file__).parent / "pe_t2i_system_prompt.txt").read_text(encoding="utf-8").strip()
PE_I2I = (Path(__file__).parent / "pe_i2i_system_prompt.txt").read_text(encoding="utf-8").strip()
ENHANCE_MAX_NEW_TOKENS = 512  # the T2I rewriter's paragraph runs 220-330 tokens; edits 60-110
CJK = re.compile(r"[\u4e00-\u9fff]")


def enhance_prompt(prompt, images, ratio_name):
    """Rewrite `prompt` with the official PE system prompt; falls back to the original if the reply does not parse."""
    if images:
        # The I2I rewriter is told to answer in the instruction's language, but the base (not PE-fine-tuned)
        # encoder drifts to Chinese for English instructions; pin the language explicitly.
        user_text = prompt + ("\n（描述用中文写。）" if CJK.search(prompt) else "\n(Write the description in English.)")
    else:
        # The T2I rewriter chooses an orientation itself unless the request states a ratio; the size menu did.
        user_text = f"{prompt}\nAspect ratio: {ratio_name}"
    refs = []
    for image in images:
        width, height, _ = calculate_dimensions(1024 * 1024, image.size[0] / image.size[1])
        refs.append(image.resize((width, height), Image.LANCZOS))
    messages = [
        {"role": "system", "content": [{"type": "text", "text": PE_I2I if images else PE_T2I}]},
        {"role": "user", "content": [{"type": "image"} for _ in refs] + [{"type": "text", "text": user_text}]},
    ]
    text = pipe.processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    inputs = pipe.processor(text=[text], images=refs or None, return_tensors="pt").to("cuda")
    with torch.inference_mode():
        tokens = pipe.text_encoder.generate(**inputs, max_new_tokens=ENHANCE_MAX_NEW_TOKENS, do_sample=False)
    reply = pipe.processor.batch_decode(tokens[:, inputs["input_ids"].shape[1] :], skip_special_tokens=True)[0]
    match = re.search(r"\{.*\}", reply, re.S)
    try:
        rewritten = json.loads(match.group())["rewritten_prompt"] if match else None
    except (ValueError, KeyError, TypeError):
        rewritten = None
    return rewritten.strip() if isinstance(rewritten, str) and rewritten.strip() else prompt


@gpu
def generate(prompt, image_1, image_2, image_3, size_label, seed, randomize_seed, enhance=True, steps=STEPS):
    steps = min(max(int(steps), 3), 8)  # 6 (RAW_NODES) is the validated schedule; 3-8 are accepted, fewer or more degrade quickly
    images = [image for image in (image_1, image_2, image_3) if image is not None]
    # unchecked "Randomize seed" always honours the box, so the UI never says one thing and does another
    seed = random.randint(0, MAX_SEED) if (randomize_seed or seed is None) else max(int(seed), 0)
    size = SIZES.get(size_label)
    width, height = size if size else (None, None)
    # Belt and braces: an editing call must land on one of the editing buckets even if the caller
    # reaches generate() with a text-to-image size (stale dropdown, gr.Examples, API client). Test
    # membership in the offered dims rather than a raw area threshold: calculate_dimensions(1536**2,
    # 4/3) is 1760x1344 = 2_365_440 px, slightly *above* 1536**2, so an area test would flag two
    # sizes the editing menu itself offers. Snap to the 1536² bucket of the nearest named ratio.
    clamped = ""
    if images and width and (width, height) not in EDIT_DIMS:
        ratio = min(RATIOS, key=lambda name_ratio: abs(name_ratio[1] - width / height))[1]
        width, height, _ = calculate_dimensions(1536 * 1536, ratio)
        clamped = f" · clamped to {width}×{height} (editing caps at the 1536² bucket)"
    used_prompt, enhance_note = prompt, ""
    if enhance and prompt.strip():
        started = time.perf_counter()
        used_prompt = enhance_prompt(prompt, images, size_label.split(" · ")[0])
        enhance_note = f" · enhance {time.perf_counter() - started:.1f}s"
        if used_prompt == prompt:
            enhance_note += " (rewrite failed to parse, original prompt used)"
    # No VAE tiling anywhere (xlarge card): tiling the reference encode wrecks edits, and tiled decodes are not the
    # validated pipeline either.
    generator = torch.Generator(device="cuda").manual_seed(seed)
    started = time.perf_counter()
    result = pipe(
        prompt=used_prompt,
        image=images or None,
        width=width,
        height=height,
        num_inference_steps=steps,
        sigmas=raw_nodes(steps),  # extra steps go to the high-noise segment, see raw_nodes()
        true_cfg_scale=1.0,  # the student is distilled without classifier-free guidance; keep it off
        output_resolution=1024,  # condition images are encoded at 1024-area, as in distillation
        generator=generator,
    ).images[0]
    elapsed = time.perf_counter() - started
    return result, f"seed `{seed}` · {result.width}×{result.height} · {steps} steps · {elapsed:.2f}s · {STUDENT_TAG}{enhance_note}{clamped}", used_prompt


def refresh_sizes(image_1, image_2, image_3, current):
    """Text-to-image and editing were distilled on different target buckets, so the menu follows the mode."""
    choices = EDIT_CHOICES if any(image is not None for image in (image_1, image_2, image_3)) else T2I_CHOICES
    return gr.update(choices=choices, value=current if current in choices else choices[0])


with gr.Blocks(title="Viggle Turbo v0.2.1 · Qwen-Image-2.1 6-step") as demo:
    gr.Markdown(
        "# Viggle Turbo v0.2.1 — 6-step Qwen-Image-2.1\n"
        "A DMD-distilled student of **Qwen-Image-2.1** that generates and edits in **6 sampling steps**, with no "
        "classifier-free guidance: about **5× faster** than the 40-step base model and very competitive with it in quality, "
        "and on some prompts we prefer its output. The **Comparison** tab has 33 examples of the official Qwen Space side by "
        "side, turbo vs base, same seed. Complicated edits (multi-reference composition, face swaps, identity-preserving "
        "edits) can still fall short of the base model.\n\n"
        "Leave the reference images empty for text-to-image; add one to three of them to edit, compose or transfer style. "
        "The size menu switches to the editing buckets as soon as a reference is attached. "
        "**v0.2.1 (2026-09-24):** the step-700 checkpoint of the v0.2 run on the 6-step schedule — intra-prompt diversity "
        "0.98× the base model, 0% composition drift; earlier versions and the numbers are on the model card. "
        f"Weights: **{STUDENT_TAG}** from [{STUDENT_REPO}](https://huggingface.co/{STUDENT_REPO})."
    )
    with gr.Tab("Generate"):
        with gr.Row():
            with gr.Column(scale=3):
                prompt = gr.Textbox(label="Prompt", lines=3, placeholder="Describe the image, or the edit to apply to the references.")
                with gr.Row():
                    image_1 = gr.Image(label="Reference 1 (optional)", type="pil", image_mode="RGB", height=200)
                    image_2 = gr.Image(label="Reference 2 (optional)", type="pil", image_mode="RGB", height=200)
                    image_3 = gr.Image(label="Reference 3 (optional)", type="pil", image_mode="RGB", height=200)
                with gr.Row():
                    # allow_custom_value: gradio validates API calls against the *initial* choices (the text-to-image
                    # menu), which would reject every editing size sent through /generate; generate() snaps
                    # anything off-menu itself.
                    size_label = gr.Dropdown(label="Output size", choices=T2I_CHOICES, value=T2I_CHOICES[0], allow_custom_value=True)
                    steps = gr.Slider(label="Steps", minimum=3, maximum=8, step=1, value=STEPS,
                                      info="6 is the validated schedule. Extra steps are always added at the high-noise end (7 also validated, 8 untested); fewer steps drop them again (4 = training nodes, 3 = uniform spacing, off-contract).")
                with gr.Row():
                    seed = gr.Number(label="Seed", value=0, precision=0, interactive=False)
                    randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
                enhance = gr.Checkbox(
                    label="Enhance prompt",
                    value=True,
                    info="Rewrite the prompt with the official Qwen-Image prompt-enhancement instructions "
                    "(runs on the built-in Qwen3-VL text encoder, +4–15 s). The text actually sent is shown under the result.",
                )
                run = gr.Button("Generate", variant="primary")
            with gr.Column(scale=4):
                output_image = gr.Image(label="Result", type="pil", height=560)
                info = gr.Markdown()
                used_prompt = gr.Textbox(label="Prompt sent to the model", lines=4, interactive=False, show_copy_button=True)

        def example_details(prompt, image_1, image_2, image_3, result):
            """Runs on an example click (no GPU): fills the status line and the prompt stored with the row, and sets the
            steps / size / enhancement controls to what produced it (rows differ; Run then reproduces the result). It also sets
            the size menu's mode itself: loading an example does not fire the reference images' .input listeners."""
            row = next(row for row in EXAMPLES if row["prompt"] == prompt)
            choices = EDIT_CHOICES if any(row["refs"]) else T2I_CHOICES
            return row["info"], row["used_prompt"], row["steps"], gr.update(choices=choices, value=row["size_label"]), row["enhance"]

        if EXAMPLES:
            gr.Examples(
                examples=[
                    [row["prompt"], *[str(EXAMPLES_DIR / ref) if ref else None for ref in row["refs"]], str(EXAMPLES_DIR / row["result"])]
                    for row in EXAMPLES
                ],
                inputs=[prompt, image_1, image_2, image_3, output_image],
                fn=example_details,
                outputs=[info, used_prompt, steps, size_label, enhance],
                run_on_click=True,
                examples_per_page=12,
                label="Examples · results pre-rendered by this model — click a row to load it (6 steps, prompt enhancement on, unless the status line says otherwise)",
            )

    with gr.Tab("Comparison"):
        compare.render()

    gr.Markdown(
        "---\n"
        "Model: [Viggle/Qwen-Image-2.1-viggle-turbo](https://huggingface.co/Viggle/Qwen-Image-2.1-viggle-turbo) · "
        "Built with Qwen — distilled from [Qwen/Qwen-Image-2.1](https://huggingface.co/Qwen/Qwen-Image-2.1), "
        "which is released under the "
        "[Qwen RESEARCH LICENSE AGREEMENT](https://huggingface.co/Qwen/Qwen-Image-2.1/blob/main/LICENSE) "
        "(**non-commercial: research or evaluation purposes only**). This demo inherits that restriction."
    )

    # .input (user upload / clear) and not .change: .change also fires when an example loads its references, and that
    # refresh then raced example_details and reset the size to the first menu entry
    for image_input in (image_1, image_2, image_3):
        image_input.input(
            refresh_sizes,
            inputs=[image_1, image_2, image_3, size_label],
            outputs=size_label,
        )
    randomize_seed.change(lambda on: gr.update(interactive=not on), inputs=randomize_seed, outputs=seed)
    run.click(
        generate,
        inputs=[prompt, image_1, image_2, image_3, size_label, seed, randomize_seed, enhance, steps],
        outputs=[output_image, info, used_prompt],
    )

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