app.py

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

import os

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

from PIL import Image
from diffusers import QwenImageEditPlusPipeline, QwenImageTransformer2DModel

BASE_REPO = "thornmaze/qwen-image-edit-2511"
BASE_REV = "1f67f908c97e2b11b235a9949632b0b1b4e5274c"
TRANSFORMER_REPO = "thornmaze/qwen-image-edit-rapid-aio-v23"
TRANSFORMER_REV = "11202b2bacd38334ff711f4f2db7aae34e48c949"

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)

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


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


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,
    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
    n_inputs = 0
    if images:
        try:
            n_inputs = len(images)
        except Exception:
            n_inputs = 1
    steps = max(1, int(num_inference_steps))
    res_scale = ((h * w) / (1024 * 1024)) ** 1.3
    estimate = int(8 + n_inputs * 1.0 + steps * 3.0 * res_scale)
    return max(20, min(120, estimate))


# 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,
    num_images_per_prompt=1,
    progress=gr.Progress(track_tqdm=True),
):
    """Run one image edit.

    `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.
    """
    negative_prompt = " "

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

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

    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 height == 256 and width == 256:
        height, width = None, None

    with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
        image = pipe(
            image=pil_images if len(pil_images) > 0 else None,
            prompt=prompt,
            height=height,
            width=width,
            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,
        ).images

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


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

        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,
        inputs=[
            input_images,
            prompt,
            seed,
            randomize_seed,
            true_guidance_scale,
            num_inference_steps,
            height,
            width,
            rewrite_prompt,
            zerogpu_budget,
        ],
        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)