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)