Blame for app.py

import gradio as gr
import numpy as np
import random
import torch
import spaces

import os
import tempfile

os.environ.setdefault("HF_HUB_DISABLE_TELEMETRY", "1")
os.environ.setdefault("DO_NOT_TRACK", "1")
os.environ.setdefault("GRADIO_ANALYTICS_ENABLED", "False")

import torch.nn.functional as F
from PIL import Image, ImageFilter
from diffusers import AutoencoderKLWan, QwenImageEditPlusPipeline, QwenImageTransformer2DModel
from diffusers.pipelines.qwenimage.pipeline_qwenimage_edit_plus import (
    VAE_IMAGE_SIZE,
    calculate_dimensions,
)

BASE_REPO = "tilnex/qwen-image-edit-2511"
BASE_REV = "fd16f1a388068cbb0cb5b2632b3624a60731ea28"
TRANSFORMER_REPO = "tilnex/qwen-image-edit-rapid-aio-v19"
TRANSFORMER_REV = "f795ea4537951c3ebf72fcdf65738035fc92bb60"
# Retrained Wan2.1 decoder that emits 12 channels and pixel-shuffles to 2x. Qwen froze Wan's
# encoder and fine-tuned only the image decoder, so the two share one latent space exactly —
# same z_dim and the same 16 latents_mean/latents_std values. That is what makes the swap legal.
UPSCALE_VAE_REPO = "tilnex/wan21-vae-upscale2x"
UPSCALE_VAE_REV = "bf80e2527d22e51b007431a364969c7fc64b08be"

dtype = torch.bfloat16
device = "cuda" if torch.cuda.is_available() else "cpu"

# No device_map: accelerate would place weights directly and bypass ZeroGPU's module-scope
# .to("cuda") interception, so they never reach the disk offload "pack".
transformer = QwenImageTransformer2DModel.from_pretrained(
    TRANSFORMER_REPO,
    revision=TRANSFORMER_REV,
    torch_dtype=dtype,
)
# Qwen-Image-Edit-2511 was trained with zero_cond_t=True (reference images get timestep=0
# modulation). The config.json in the mirror carries it; this line is a safety net.
transformer.config.zero_cond_t = True

pipe = QwenImageEditPlusPipeline.from_pretrained(
    BASE_REPO,
    revision=BASE_REV,
    transformer=transformer,
    torch_dtype=dtype,
).to(device)

# Decoder only — the pipeline's own VAE still encodes the input image, so both stay loaded.
upscale_vae = AutoencoderKLWan.from_pretrained(
    UPSCALE_VAE_REPO,
    revision=UPSCALE_VAE_REV,
    torch_dtype=dtype,
).to(device)

MAX_SEED = np.iinfo(np.int32).max

# Edit-style LoRA adapters, self-hosted on the owner account and preloaded at build (README) exactly like
# the base + transformer. Loaded at import in the main process, where the network and HF cache are
# available; the @spaces.GPU task only toggles set_adapters (its sandbox can neither download nor
# write the cache). NO_ADAPTER = the baked pipeline, unchanged behaviour.
NO_ADAPTER = "None"

ADAPTER_SPECS = {
    "Photo-to-Anime": {
        "repo": "tilnex/qwen-image-edit-2509-photo-to-anime",
        "revision": "c356b2a1d81bc91e2d60b87c4828d6ab770b7ba4",
        "weights": "Qwen-Image-Edit-2509-Photo-to-Anime_000001000.safetensors",
        "adapter_name": "photo-to-anime",
    },
    "Multiple-Angles": {
        "repo": "tilnex/qie-loras",
        "revision": "cf6abb05c1775075e74830855b7c3eeff81fa8be",
        "weights": "multiple-angles.safetensors",
        "adapter_name": "multiple-angles",
    },
    "Anime-V2": {
        "repo": "tilnex/qie-loras",
        "revision": "cf6abb05c1775075e74830855b7c3eeff81fa8be",
        "weights": "anime-v2.safetensors",
        "adapter_name": "anime-v2",
    },
    "Light-Migration": {
        "repo": "tilnex/qie-loras",
        "revision": "cf6abb05c1775075e74830855b7c3eeff81fa8be",
        "weights": "light-migration.safetensors",
        "adapter_name": "light-migration",
    },
    "Upscaler": {
        "repo": "tilnex/qie-loras",
        "revision": "cf6abb05c1775075e74830855b7c3eeff81fa8be",
        "weights": "upscaler.safetensors",
        "adapter_name": "upscaler",
    },
    "Style-Transfer": {
        "repo": "tilnex/qie-loras",
        "revision": "cf6abb05c1775075e74830855b7c3eeff81fa8be",
        "weights": "style-transfer.safetensors",
        "adapter_name": "style-transfer",
    },
    "Manga-Tone": {
        "repo": "tilnex/qie-loras",
        "revision": "cf6abb05c1775075e74830855b7c3eeff81fa8be",
        "weights": "manga-tone.safetensors",
        "adapter_name": "manga-tone",
    },
    "Anything2Real": {
        "repo": "tilnex/qie-loras",
        "revision": "cf6abb05c1775075e74830855b7c3eeff81fa8be",
        "weights": "anything2real.safetensors",
        "adapter_name": "anything2real",
    },
    "Polaroid-Photo": {
        "repo": "tilnex/qie-loras",
        "revision": "cf6abb05c1775075e74830855b7c3eeff81fa8be",
        "weights": "polaroid-photo.safetensors",
        "adapter_name": "polaroid-photo",
    },
    "Midnight-Noir-Eyes-Spotlight": {
        "repo": "tilnex/qie-loras",
        "revision": "cf6abb05c1775075e74830855b7c3eeff81fa8be",
        "weights": "midnight-noir-eyes-spotlight.safetensors",
        "adapter_name": "midnight-noir-eyes-spotlight",
    },
    "Hyper-Realistic-Portrait": {
        "repo": "tilnex/qie-loras",
        "revision": "cf6abb05c1775075e74830855b7c3eeff81fa8be",
        "weights": "hyper-realistic-portrait.safetensors",
        "adapter_name": "hyper-realistic-portrait",
    },
    "Ultra-Realistic-Portrait": {
        "repo": "tilnex/qie-loras",
        "revision": "cf6abb05c1775075e74830855b7c3eeff81fa8be",
        "weights": "ultra-realistic-portrait.safetensors",
        "adapter_name": "ultra-realistic-portrait",
    },
    "Pixar-Inspired-3D": {
        "repo": "tilnex/qie-loras",
        "revision": "cf6abb05c1775075e74830855b7c3eeff81fa8be",
        "weights": "pixar-inspired-3d.safetensors",
        "adapter_name": "pixar-inspired-3d",
    },
    "Noir-Comic-Book": {
        "repo": "tilnex/qie-loras",
        "revision": "cf6abb05c1775075e74830855b7c3eeff81fa8be",
        "weights": "noir-comic-book.safetensors",
        "adapter_name": "noir-comic-book",
    },
    "Any-light": {
        "repo": "tilnex/qie-loras",
        "revision": "cf6abb05c1775075e74830855b7c3eeff81fa8be",
        "weights": "any-light.safetensors",
        "adapter_name": "any-light",
    },
    "Studio-DeLight": {
        "repo": "tilnex/qie-loras",
        "revision": "cf6abb05c1775075e74830855b7c3eeff81fa8be",
        "weights": "studio-delight.safetensors",
        "adapter_name": "studio-delight",
    },
    "4c98877c": {
        "repo": "tilnex/qie-loras",
        "revision": "cf6abb05c1775075e74830855b7c3eeff81fa8be",
        "weights": "4c98877c.safetensors",
        "adapter_name": "4c98877c",
    },
    "b7862b96": {
        "repo": "tilnex/qie-loras",
        "revision": "cf6abb05c1775075e74830855b7c3eeff81fa8be",
        "weights": "b7862b96.safetensors",
        "adapter_name": "b7862b96",
    },
    "42dc0cf4": {
        "repo": "tilnex/qie-loras",
        "revision": "cf6abb05c1775075e74830855b7c3eeff81fa8be",
        "weights": "42dc0cf4.safetensors",
        "adapter_name": "42dc0cf4",
    },
    "3bdf5c5a": {
        "repo": "tilnex/qie-loras",
        "revision": "cf6abb05c1775075e74830855b7c3eeff81fa8be",
        "weights": "3bdf5c5a.safetensors",
        "adapter_name": "3bdf5c5a",
    },
    "Lowres-Fix": {
        "repo": "tilnex/qie-loras",
        "revision": "cf6abb05c1775075e74830855b7c3eeff81fa8be",
        "weights": "lowres-fix.safetensors",
        "adapter_name": "lowres-fix",
    },
    "Fal-Multiple-Angles": {
        "repo": "tilnex/qie-loras",
        "revision": "cf6abb05c1775075e74830855b7c3eeff81fa8be",
        "weights": "fal-multiple-angles.safetensors",
        "adapter_name": "fal-multiple-angles",
    },
    "Unblur-Anything": {
        "repo": "tilnex/qie-loras",
        "revision": "cf6abb05c1775075e74830855b7c3eeff81fa8be",
        "weights": "unblur-anything.safetensors",
        "adapter_name": "unblur-anything",
    },
    "Cinematic-FlatLog": {
        "repo": "tilnex/qie-loras",
        "revision": "cf6abb05c1775075e74830855b7c3eeff81fa8be",
        "weights": "cinematic-flatlog.safetensors",
        "adapter_name": "cinematic-flatlog",
    },
}

for _spec in ADAPTER_SPECS.values():
    # The revision is pinned here as well as in the card: the card's preload seeds a cache that is
    # READ-ONLY at runtime, so an unpinned load asking for a newer file would try to download into
    # it and die with a permission error instead of a missing-file one.
    pipe.load_lora_weights(
        _spec["repo"],
        revision=_spec["revision"],
        weight_name=_spec["weights"],
        adapter_name=_spec["adapter_name"],
    )
pipe.disable_lora()  # default off — _apply_adapters turns on exactly what the call names


# How the latents become pixels. The dot/halftone grid on skin is produced by the 2511 decoder at
# its own 8 px stride, so every option here is about that stride and nothing else.
MODE_AS_IS = "as-is"  # exactly what the pool Spaces do today
MODE_RENDER2X = "render2x"  # render at twice the area, Lanczos back down
MODE_VAE2X = "vae2x"  # decode 2x with the retrained decoder, Lanczos 0.5
DECODE_MODES = (MODE_AS_IS, MODE_RENDER2X, MODE_VAE2X)

# The pipeline floors both sides to vae_scale_factor * 2 after computing the target.
_MULTIPLE_OF = 16


def _floor16(width: int, height: int):
    return width // _MULTIPLE_OF * _MULTIPLE_OF, height // _MULTIPLE_OF * _MULTIPLE_OF


def _sizes(pil_images, height, width, decode_mode):
    """(render, deliver) rasters for this call.

    `deliver` reproduces the pipeline's own arithmetic — target area at the LAST input's aspect,
    each side rounded to 32 then floored to 16 — so `as-is` stays byte-identical to a call that
    passes height/width as None. `render` is what the transformer is actually asked for.
    """
    if height and width and int(height) > 256 and int(width) > 256:
        deliver = _floor16(int(width), int(height))
    elif pil_images:
        src_w, src_h = pil_images[-1].size
        deliver = _floor16(*calculate_dimensions(VAE_IMAGE_SIZE, src_w / src_h))
    else:
        deliver = (1024, 1024)
    if decode_mode == MODE_RENDER2X:
        ratio = deliver[0] / deliver[1]
        return _floor16(*calculate_dimensions(VAE_IMAGE_SIZE * 2, ratio)), deliver
    return deliver, deliver


def _resolve_adapters(lora_adapter):
    """The specs to stack for this call, in the caller's order.

    Accepts a LIST (the multiselect dropdown and the pool both send one) or a bare string: the
    Space and the backend deploy separately, so a caller still on the single-value contract has
    to keep working through the window between the two rollouts."""
    names = [lora_adapter] if isinstance(lora_adapter, str) else list(lora_adapter or ())
    seen, specs = set(), []
    for name in names:
        if not name or name == NO_ADAPTER or name in seen:
            continue
        spec = ADAPTER_SPECS.get(name)
        if spec is not None:
            seen.add(name)
            specs.append(spec)
    return specs


def _apply_adapters(specs) -> None:
    """Toggle the LoRAs for this call — every adapter is already loaded at import.

    Stacked adapters all carry weight 1.0 and diffusers SUMS their deltas, so three strong ones
    over-cook the edit. No cap here: which combinations are usable is measured on real photos,
    not guessed in the handler."""
    if not specs:
        pipe.disable_lora()
        return
    pipe.enable_lora()
    names = [s["adapter_name"] for s in specs]
    pipe.set_adapters(names, adapter_weights=[1.0] * len(names))


@torch.no_grad()
def _decode_2x(packed_latents, render):
    """Decode packed latents with the 2x decoder and return PIL images at twice `render`.

    Repeats the pipeline's own unpack + denormalise (it only runs them on the non-latent path),
    then takes the 12-channel output the retrained decoder emits and pixel-shuffles it to RGB at
    2x. Deliberately NOT under autocast: this is the one place a different decode dtype could ever
    be tried, and autocast would silently pin it to bfloat16. `no_grad` because the pipeline's own
    decode runs inside its inference context and this one does not — without it the result carries
    a grad graph and `.numpy()` refuses.
    """
    latents = pipe._unpack_latents(packed_latents, render[1], render[0], pipe.vae_scale_factor)
    latents = latents.to(upscale_vae.dtype)
    z_dim = pipe.vae.config.z_dim
    mean = (
        torch.tensor(pipe.vae.config.latents_mean)
        .view(1, z_dim, 1, 1, 1)
        .to(latents.device, latents.dtype)
    )
    inv_std = (
        (1.0 / torch.tensor(pipe.vae.config.latents_std))
        .view(1, z_dim, 1, 1, 1)
        .to(latents.device, latents.dtype)
    )
    latents = latents / inv_std + mean
    decoded = upscale_vae.decode(latents, return_dict=False)[0]  # [B, 12, 1, H, W]
    rgb = F.pixel_shuffle(decoded.squeeze(2), 2)  # [B, 3, 2H, 2W]
    rgb = (rgb / 2 + 0.5).clamp(0, 1)
    frames = (rgb.permute(0, 2, 3, 1).float().cpu().numpy() * 255).round().astype("uint8")
    return [Image.fromarray(frame) for frame in frames]


def use_output_as_input(output_images):
    """Convert output images to input format for the gallery"""
    if output_images is None or len(output_images) == 0:
        return []
    return output_images


# Warm generation time measured 31.08 on the pool at 4 steps: 4.2 s at 1.06 MP, 6.5 s at 2.17 MP,
# i.e. ~2 s fixed plus ~0.52 s per megapixel-step. The declaration doubles that and adds the cold
# restore, which runs INSIDE the declared window; the 20 s floor is what that restore needs, and
# over-declaring costs pool availability (an account below the threshold is skipped entirely).
_FIXED_SEC = 2.0
_SEC_PER_MP_STEP = 0.52
_RESTORE_SEC = 10.0


def get_edit_duration(
    images,
    prompt,
    seed=42,
    randomize_seed=False,
    true_guidance_scale=1.0,
    num_inference_steps=4,
    height=None,
    width=None,
    rewrite_prompt=True,
    zerogpu_budget=0,
    lora_adapter=NO_ADAPTER,
    decode_mode=MODE_VAE2X,
    num_images_per_prompt=1,
    progress=None,
):
    if zerogpu_budget and int(zerogpu_budget) > 0:
        return max(20, min(120, int(zerogpu_budget)))
    h = int(height) if height and int(height) > 256 else 1024
    w = int(width) if width and int(width) > 256 else 1024
    megapixels = (h * w) / (1024 * 1024)
    if decode_mode == MODE_RENDER2X:
        megapixels *= 2
    steps = max(1, int(num_inference_steps))
    compute = _FIXED_SEC + _SEC_PER_MP_STEP * megapixels * steps
    if decode_mode == MODE_VAE2X:
        compute += 2.0  # the 2x decode, measured against the 1x one
    return max(20, min(120, int(_RESTORE_SEC + 2 * compute)))


# size="xlarge" is MANDATORY: the packed pipeline is 57.7GB and `large` gives 48GB.
# Without it a freshly created Space fails every GPU call with a bodiless error.
@spaces.GPU(duration=get_edit_duration, size="xlarge")
def infer(
    images,
    prompt,
    seed=42,
    randomize_seed=False,
    true_guidance_scale=1.0,
    num_inference_steps=4,
    height=None,
    width=None,
    rewrite_prompt=True,
    zerogpu_budget=0,
    lora_adapter=NO_ADAPTER,
    decode_mode=MODE_VAE2X,
    num_images_per_prompt=1,
    progress=gr.Progress(track_tqdm=True),
):
    """Run one image edit. Returns the images BEFORE any downscale, plus the delivery raster.

    `rewrite_prompt` is accepted but unused — the caller's positional argument array is a
    fixed contract; dropping the slot would shift every argument after it. `lora_adapter` and
    `decode_mode` are appended AFTER that legacy contract, so an 11-arg caller lands on
    MODE_AS_IS and behaves exactly like the pool Spaces.
    """
    negative_prompt = " "

    if randomize_seed:
        seed = random.randint(0, MAX_SEED)

    generator = torch.Generator(device=device).manual_seed(seed)

    _apply_adapters(_resolve_adapters(lora_adapter))

    pil_images = []
    if images is not None:
        for item in images:
            try:
                if isinstance(item[0], Image.Image):
                    pil_images.append(item[0].convert("RGB"))
                elif isinstance(item[0], str):
                    pil_images.append(Image.open(item[0]).convert("RGB"))
                elif hasattr(item, "name"):
                    pil_images.append(Image.open(item.name).convert("RGB"))
            except Exception:
                continue

    if decode_mode not in DECODE_MODES:
        decode_mode = MODE_VAE2X
    render, deliver = _sizes(pil_images, height, width, decode_mode)

    # Autocast wraps the TRANSFORMER only. It is a no-op on weights already in bfloat16, but
    # leaving it around the decode would pin that decode's dtype too.
    with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
        result = pipe(
            image=pil_images if len(pil_images) > 0 else None,
            prompt=prompt,
            height=render[1],
            width=render[0],
            negative_prompt=negative_prompt,
            num_inference_steps=num_inference_steps,
            generator=generator,
            true_cfg_scale=true_guidance_scale,
            num_images_per_prompt=num_images_per_prompt,
            output_type="latent" if decode_mode == MODE_VAE2X else "pil",
        ).images

    image = _decode_2x(result, render) if decode_mode == MODE_VAE2X else result

    return image, seed, gr.update(visible=True), deliver


# Gradio re-encodes a returned PIL image with its own default (WebP quality 80) — a second lossy
# pass over what the model produced, on top of nothing. `format=` only governs numpy/PIL returns, so
# handing back a PATH is what keeps that pass out of the picture entirely.
#
# 98 rather than 95: measured 31.08 on a delivered frame, q95 -> q98 costs 45 KB and takes the worst
# error on smooth skin from 10/255 to 7. Above q98 the ladder is flat (the codec subsamples chroma at
# every quality, so ~47 dB is its ceiling) and lossless is 674 KB — at 30-day retention and 2.5k
# photos a day that is 38-50 GB standing against 57 GB free.
RESULT_WEBP_QUALITY = 98


def infer_to_files(
    images,
    prompt,
    seed=42,
    randomize_seed=False,
    true_guidance_scale=1.0,
    num_inference_steps=4,
    height=None,
    width=None,
    rewrite_prompt=True,
    zerogpu_budget=0,
    lora_adapter=NO_ADAPTER,
    decode_mode=MODE_VAE2X,
    num_images_per_prompt=1,
    progress=gr.Progress(track_tqdm=True),
):
    """Run the edit, downscale to the delivery raster, encode here and return file paths.

    Deliberately NOT inside `infer`: that function holds the ZeroGPU allocation, and a resize or
    an encode there would be billed as GPU seconds. The downscale IS the halftone fix — the grid
    lives at the decoder's own stride, and resampling below that stride is what removes it.
    """
    images_out, used_seed, btn, deliver = infer(
        images,
        prompt,
        seed,
        randomize_seed,
        true_guidance_scale,
        num_inference_steps,
        height,
        width,
        rewrite_prompt,
        zerogpu_budget,
        lora_adapter,
        decode_mode,
        num_images_per_prompt,
        progress,
    )
    paths = []
    for image in images_out:
        if image.size != tuple(deliver):
            if decode_mode == MODE_VAE2X:
                # The decoder's author asks for a slight blur before HIS exact 0.5 downscale,
                # against aliasing. Measured as a no-op at render2x's 0.707 ratio, so it is applied
                # only where he specified it.
                image = image.filter(ImageFilter.GaussianBlur(0.25))
            image = image.resize(tuple(deliver), Image.Resampling.LANCZOS)
        fd, path = tempfile.mkstemp(suffix=".webp")
        os.close(fd)
        image.save(path, format="WEBP", quality=RESULT_WEBP_QUALITY, method=4)
        paths.append(path)
    return paths, used_seed, btn


css = """
#col-container {
    margin: 0 auto;
    max-width: 1024px;
}
"""

with gr.Blocks(analytics_enabled=False) as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown("### Image Lab")
        with gr.Row():
            with gr.Column():
                input_images = gr.Gallery(
                    label="Input Images", show_label=False, type="pil", interactive=True
                )
            with gr.Column():
                result = gr.Gallery(label="Result", show_label=False, type="pil", interactive=False)
                use_output_btn = gr.Button(
                    "↗️ Use as input", variant="secondary", size="sm", visible=False
                )

        with gr.Row():
            prompt = gr.Text(
                label="Prompt",
                show_label=False,
                placeholder="describe the edit instruction",
                container=False,
            )
            run_button = gr.Button("Edit!", variant="primary")

        lora_adapter = gr.Dropdown(
            label="Choose Editing Styles",
            choices=list(ADAPTER_SPECS),
            value=[],
            multiselect=True,
        )

        with gr.Accordion("Advanced Settings", open=False):
            seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0)
            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
            with gr.Row():
                true_guidance_scale = gr.Slider(
                    label="True guidance scale", minimum=1.0, maximum=10.0, step=0.1, value=1.0
                )
                num_inference_steps = gr.Slider(
                    label="Number of inference steps", minimum=1, maximum=40, step=1, value=4
                )
                height = gr.Slider(label="Height", minimum=256, maximum=2048, step=8, value=None)
                width = gr.Slider(label="Width", minimum=256, maximum=2048, step=8, value=None)
                rewrite_prompt = gr.Checkbox(label="Rewrite prompt (inactive)", value=False)
            zerogpu_budget = gr.Slider(
                label="ZeroGPU max duration (0 = auto)", minimum=0, maximum=120, step=5, value=0
            )

    gr.on(
        triggers=[run_button.click, prompt.submit],
        fn=infer_to_files,
        # The route name is derived from the function name unless pinned, and the bot calls
        # `/gradio_api/call/infer` — renaming the handler would 404 every generation.
        api_name="infer",
        # Exactly the pool's 11 slots. `decode_mode` is a handler default, never an input: gradio
        # rejects a call carrying FEWER values than the handler has inputs.
        inputs=[
            input_images,
            prompt,
            seed,
            randomize_seed,
            true_guidance_scale,
            num_inference_steps,
            height,
            width,
            rewrite_prompt,
            zerogpu_budget,
            lora_adapter,
        ],
        outputs=[result, seed, use_output_btn],
    )

    use_output_btn.click(fn=use_output_as_input, inputs=[result], outputs=[input_images])

if __name__ == "__main__":
    # ssr_mode is Gradio 6's experimental server-side render node; with it on the Space
    # intermittently answers 500 at the edge while the app process is healthy.
    demo.launch(css=css, show_error=True, ssr_mode=False)