import gradio as gr import numpy as np import random import torch import spaces from PIL import Image from torchao.quantization import Float8DynamicActivationFloat8WeightConfig from torchao.quantization import quantize_ from qwenimage.pipeline_qwenimage_edit_plus import QwenImageEditPlusPipeline from qwenimage.transformer_qwenimage import QwenImageTransformer2DModel from qwenimage.qwen_fa3_processor import QwenDoubleStreamAttnProcessorFA3 import time def update_history(new_images, history): """Updates the history gallery with the new images.""" time.sleep(0.5) # Small delay to ensure images are ready if history is None: history = [] if new_images is not None and len(new_images) > 0: if not isinstance(history, list): history = list(history) if history else [] for img in new_images: history.insert(0, img) history = history[:20] # Keep only last 20 images return history def use_history_as_input(evt: gr.SelectData): """Sets the selected history image into the Image 1 slot.""" if evt.value is not None: # gr.Image with type='filepath' accepts a path directly. return gr.update(value=evt.value) return gr.update() # --- Model Loading --- dtype = torch.bfloat16 device = "cuda" if torch.cuda.is_available() else "cpu" BASE_MODEL_ID = "Qwen/Qwen-Image-Edit-2511" BASE_MODEL_REVISION = "6f3ccc0b56e431dc6a0c2b2039706d7d26f22cb9" ACCELERATED_TRANSFORMER_ID = "Sneak-Moose/Qwen-Rapid-AIO-v18-NSFW-diffusers" ACCELERATED_TRANSFORMER_REVISION = "5641245ab83ffd498c986485eb8c3e9f6f3f2184" # This UI is tuned for four-step inference and must use the matching accelerated # transformer. Do not silently fall back to the standard transformer: doing so # produces misleading low-quality output while presenting the app as healthy. try: transformer = QwenImageTransformer2DModel.from_pretrained( ACCELERATED_TRANSFORMER_ID, subfolder="transformer", revision=ACCELERATED_TRANSFORMER_REVISION, torch_dtype=dtype, device_map="cuda" if torch.cuda.is_available() else None, ) except Exception as exc: raise RuntimeError( "The pinned four-step accelerated transformer could not be loaded. " "Generation is disabled rather than silently using an incompatible fallback." ) from exc pipe = QwenImageEditPlusPipeline.from_pretrained( BASE_MODEL_ID, revision=BASE_MODEL_REVISION, transformer=transformer, torch_dtype=dtype, ) del transformer if torch.cuda.is_available(): torch.cuda.empty_cache() # A full BF16 pipeline is roughly 58 GB and cannot be transferred into a # standard ZeroGPU allocation. Dynamic FP8 cuts the transformer/text-encoder # footprint enough to fit while preserving four-step Rapid-AIO behavior. fp8_config = Float8DynamicActivationFloat8WeightConfig() quantize_(pipe.transformer, fp8_config) try: quantize_(pipe.text_encoder, fp8_config) except Exception as exc: print( f"Text-encoder FP8 quantization unavailable ({type(exc).__name__}); " "continuing with the pinned text encoder.", flush=True, ) pipe = pipe.to(device) # Apply the upstream FA3 optimization, but keep boot resilient if kernel support # differs on duplicated Spaces or the public fallback transformer. try: pipe.transformer.__class__ = QwenImageTransformer2DModel pipe.transformer.set_attn_processor(QwenDoubleStreamAttnProcessorFA3()) except Exception as exc: print(f"FA3 attention processor unavailable; continuing with default attention: {exc}") # Load next-scene LoRA for cinematic progression # Note: This LoRA was trained on 2509, may need testing with 2511/v18 # TODO: Re-enable after testing base 2511/v18 works correctly # pipe.load_lora_weights( # "lovis93/next-scene-qwen-image-lora-2509", # weight_name="next-scene_lora-v2-3000.safetensors", # adapter_name="next-scene" # ) # pipe.set_adapters(["next-scene"], adapter_weights=[1.]) # pipe.fuse_lora(adapter_names=["next-scene"], lora_scale=1.) # pipe.unload_lora_weights() # --- Ahead-of-time compilation --- # Note: optimize_pipeline_ handles text encoder offloading internally to save memory during torch.export # DISABLED 2026-05-12: HF build pipeline force-pins spaces==0.49.3 which has a regression in # zero.torch.patching._move() — NVML assert during worker_init kills AOTI compile at startup. # Restore once HF bumps the pipeline to spaces==0.50.0+. # optimize_pipeline_(pipe, image=[Image.new("RGB", (1024, 1024)), Image.new("RGB", (1024, 1024))], prompt="prompt") # --- UI Constants and Helpers --- MAX_SEED = np.iinfo(np.int32).max MIN_INPUT_SIDE = 256 def use_output_as_input(output_images): """Move the first output image into the Image 1 slot.""" if not output_images: return gr.update() first = output_images[0] # Gallery items can be filepath strings or (filepath, label) tuples. path = first[0] if isinstance(first, (list, tuple)) else first return gr.update(value=path) # --- Main Inference Function (with hardcoded negative prompt) --- @spaces.GPU(duration=180) def infer( image_1, image_2, prompt, seed=42, randomize_seed=False, true_guidance_scale=1.0, num_inference_steps=4, height=None, width=None, num_images_per_prompt=1, progress=gr.Progress(track_tqdm=True), ): """ Generates an image using the local Qwen-Image diffusers pipeline. """ # Hardcode the negative prompt as requested negative_prompt = " " if randomize_seed: seed = random.randint(0, MAX_SEED) # Set up the generator for reproducibility generator = torch.Generator(device=device).manual_seed(seed) # Load input images into PIL Images — two optional slots. pil_images = [] for img in (image_1, image_2): if img is None: continue try: if isinstance(img, str): pil_images.append(Image.open(img).convert("RGB")) elif isinstance(img, Image.Image): pil_images.append(img.convert("RGB")) elif hasattr(img, "name"): pil_images.append(Image.open(img.name).convert("RGB")) except Exception: continue if not pil_images: raise gr.Error("Upload at least one valid input image before editing.") for img in pil_images: if min(img.size) < MIN_INPUT_SIDE: raise gr.Error( f"Input images must be at least {MIN_INPUT_SIDE}px on each side. " "Very small source images can produce tiled or grid-like artifacts." ) print( f"Starting generation: seed={seed}, steps={num_inference_steps}, " f"guidance={true_guidance_scale}, size={width}x{height}, inputs={len(pil_images)}", flush=True, ) # Generate the image images_pil = 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 # Let Gradio manage temporary result files so delete_cache can expire them. return images_pil, seed, gr.update(visible=True) # --- UI Layout --- css = """ #col-container { margin: 0 auto; max-width: 1024px; } #logo-title { text-align: center; } #logo-title img { width: 400px; } #edit_text{margin-top: -62px !important} """ with gr.Blocks(css=css, delete_cache=(3600, 86400)) as demo: with gr.Column(elem_id="col-container"): gr.HTML("""

Pro Realism Edit Studio 🎨

Rapid Edit ⚡

""") gr.Markdown(""" This demo uses [Qwen-Image-Edit-2511](https://huggingface.co/Qwen/Qwen-Image-Edit-2511) with [Phr00t's Rapid-AIO v18](https://huggingface.co/Phr00t/Qwen-Image-Edit-Rapid-AIO) accelerated transformer + [AoT compilation & FA3](https://huggingface.co/blog/zerogpu-aoti) for fast 4-step inference. Upload an image and enter your prompt to edit it. The model will use your prompt exactly as provided. """) with gr.Row(): with gr.Column(): with gr.Row(): image_1 = gr.Image(label="Image 1", type="filepath", interactive=True) image_2 = gr.Image(label="Image 2 (optional)", type="filepath", interactive=True) prompt = gr.Text( label="Prompt 🪄", show_label=True, placeholder="Enter your prompt here...", ) 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, ) with gr.Column(): result = gr.Gallery(label="Result", show_label=False, type="filepath") with gr.Row(): use_output_btn = gr.Button("↗️ Use as input", variant="secondary", size="sm", visible=False) with gr.Row(visible=False): gr.Markdown("### 📜 History") clear_history_button = gr.Button("🗑️ Clear History", size="sm", variant="stop") history_gallery = gr.Gallery( label="Click any image to use as input", interactive=False, show_label=True, visible=False ) gr.on( triggers=[run_button.click, prompt.submit], fn=infer, inputs=[ image_1, image_2, prompt, seed, randomize_seed, true_guidance_scale, num_inference_steps, ], outputs=[result, seed, use_output_btn], ).then( fn=update_history, inputs=[result, history_gallery], outputs=history_gallery, ) # Add the new event handler for the "Use Output as Input" button use_output_btn.click( fn=use_output_as_input, inputs=[result], outputs=[image_1] ) # History gallery event handlers history_gallery.select( fn=use_history_as_input, inputs=None, outputs=[image_1], ) clear_history_button.click( fn=lambda: [], inputs=None, outputs=history_gallery, ) if __name__ == "__main__": demo.launch()