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)