import os import re import gc import traceback import random from typing import Iterable, Optional import gradio as gr import numpy as np import spaces import torch from PIL import Image from transformers import AutoImageProcessor, AutoModelForDepthEstimation from huggingface_hub import hf_hub_download from safetensors.torch import load_file as safetensors_load_file from gradio.themes import Soft from gradio.themes.utils import colors, fonts, sizes # ============================================================ # Theme # ============================================================ colors.orange_red = colors.Color( name="orange_red", c50="#FFF0E5", c100="#FFE0CC", c200="#FFC299", c300="#FFA366", c400="#FF8533", c500="#FF4500", c600="#E63E00", c700="#CC3700", c800="#B33000", c900="#992900", c950="#802200", ) class OrangeRedTheme(Soft): def __init__( self, *, primary_hue: colors.Color | str = colors.gray, secondary_hue: colors.Color | str = colors.orange_red, neutral_hue: colors.Color | str = colors.slate, text_size: sizes.Size | str = sizes.text_lg, font: fonts.Font | str | Iterable[fonts.Font | str] = ( fonts.GoogleFont("Outfit"), "Arial", "sans-serif", ), font_mono: fonts.Font | str | Iterable[fonts.Font | str] = ( fonts.GoogleFont("IBM Plex Mono"), "ui-monospace", "monospace", ), ): super().__init__( primary_hue=primary_hue, secondary_hue=secondary_hue, neutral_hue=neutral_hue, text_size=text_size, font=font, font_mono=font_mono, ) super().set( background_fill_primary="*primary_50", background_fill_primary_dark="*primary_900", body_background_fill="linear-gradient(135deg, *primary_200, *primary_100)", body_background_fill_dark="linear-gradient(135deg, *primary_900, *primary_800)", button_primary_text_color="white", button_primary_text_color_hover="white", button_primary_background_fill="linear-gradient(90deg, *secondary_500, *secondary_600)", button_primary_background_fill_hover="linear-gradient(90deg, *secondary_600, *secondary_700)", button_primary_background_fill_dark="linear-gradient(90deg, *secondary_600, *secondary_700)", button_primary_background_fill_hover_dark="linear-gradient(90deg, *secondary_500, *secondary_600)", button_secondary_text_color="black", button_secondary_text_color_hover="white", button_secondary_background_fill="linear-gradient(90deg, *primary_300, *primary_300)", button_secondary_background_fill_hover="linear-gradient(90deg, *primary_400, *primary_400)", button_secondary_background_fill_dark="linear-gradient(90deg, *primary_500, *primary_600)", button_secondary_background_fill_hover_dark="linear-gradient(90deg, *primary_500, *primary_500)", slider_color="*secondary_500", slider_color_dark="*secondary_600", block_title_text_weight="600", block_border_width="3px", block_shadow="*shadow_drop_lg", button_primary_shadow="*shadow_drop_lg", button_large_padding="11px", color_accent_soft="*primary_100", block_label_background_fill="*primary_200", ) orange_red_theme = OrangeRedTheme() # ============================================================ # Device # ============================================================ device = torch.device("cuda" if torch.cuda.is_available() else "cpu") print("CUDA_VISIBLE_DEVICES=", os.environ.get("CUDA_VISIBLE_DEVICES")) print("torch.__version__ =", torch.__version__) print("torch.version.cuda =", torch.version.cuda) print("cuda available:", torch.cuda.is_available()) print("cuda device count:", torch.cuda.device_count()) if torch.cuda.is_available(): print("current device:", torch.cuda.current_device()) print("device name:", torch.cuda.get_device_name(torch.cuda.current_device())) print("Using device:", device) dtype = torch.bfloat16 MAX_SEED = np.iinfo(np.int32).max # ============================================================ # AIO version (Space variable) # ============================================================ AIO_REPO_ID = "Pr0f3ssi0n4ln00b/Phr00t-Qwen-Rapid-AIO" DEFAULT_AIO_VERSION = "v19" _VER_RE = re.compile(r"^v\d+$") _DIGITS_RE = re.compile(r"^\d+$") def _normalize_version(raw: str) -> Optional[str]: if raw is None: return None s = str(raw).strip() if not s: return None if _VER_RE.fullmatch(s): return s if _DIGITS_RE.fullmatch(s): return f"v{s}" return None _AIO_ENV_RAW = os.environ.get("AIO_VERSION", "") _AIO_ENV_NORM = _normalize_version(_AIO_ENV_RAW) AIO_VERSION = _AIO_ENV_NORM or DEFAULT_AIO_VERSION AIO_VERSION_SOURCE = "env" if _AIO_ENV_NORM else "default(v19)" print(f"AIO_VERSION (env raw) = {_AIO_ENV_RAW!r}") print(f"AIO_VERSION (normalized) = {_AIO_ENV_NORM!r}") print(f"Using AIO_VERSION = {AIO_VERSION} ({AIO_VERSION_SOURCE})") # ============================================================ # Pipeline # ============================================================ from diffusers import FlowMatchEulerDiscreteScheduler # noqa: F401 from qwenimage.pipeline_qwenimage_edit_plus import QwenImageEditPlusPipeline from qwenimage.transformer_qwenimage import QwenImageTransformer2DModel from qwenimage.qwen_fa3_processor import QwenDoubleStreamAttnProcessorFA3 def _load_pipe_with_version(version: str) -> QwenImageEditPlusPipeline: sub = f"{version}/transformer" print(f"Loading AIO transformer: {AIO_REPO_ID} / {sub}") p = QwenImageEditPlusPipeline.from_pretrained( "Qwen/Qwen-Image-Edit-2511", transformer=QwenImageTransformer2DModel.from_pretrained( AIO_REPO_ID, subfolder=sub, torch_dtype=dtype, device_map="cuda", ), torch_dtype=dtype, ).to(device) return p try: pipe = _load_pipe_with_version(AIO_VERSION) except Exception: print("❌ Failed to load requested AIO_VERSION. Falling back to v19.") print("---- exception ----") print(traceback.format_exc()) print("-------------------") AIO_VERSION = DEFAULT_AIO_VERSION AIO_VERSION_SOURCE = "fallback_to_v19" pipe = _load_pipe_with_version(AIO_VERSION) try: pipe.transformer.set_attn_processor(QwenDoubleStreamAttnProcessorFA3()) print("Flash Attention 3 Processor set successfully.") except Exception as e: print(f"Warning: Could not set FA3 processor: {e}") # ============================================================ # VAE tiling toggle (UI-controlled; OFF by default) # ============================================================ def _apply_vae_tiling(enabled: bool): """ Toggle VAE tiling on the global pipeline. This does NOT require a Space restart; it applies to the next pipe(...) call. """ try: if enabled: if hasattr(pipe, "enable_vae_tiling"): pipe.enable_vae_tiling() print("✅ VAE tiling ENABLED (per UI).") elif hasattr(pipe, "vae") and hasattr(pipe.vae, "enable_tiling"): pipe.vae.enable_tiling() print("✅ VAE tiling ENABLED via pipe.vae.enable_tiling() (per UI).") else: print("⚠️ No enable_vae_tiling()/vae.enable_tiling() found; cannot enable.") else: if hasattr(pipe, "disable_vae_tiling"): pipe.disable_vae_tiling() print("VAE tiling DISABLED (per UI).") elif hasattr(pipe, "vae") and hasattr(pipe.vae, "disable_tiling"): pipe.vae.disable_tiling() print("VAE tiling DISABLED via pipe.vae.disable_tiling() (per UI).") else: print("⚠️ No disable_vae_tiling()/vae.disable_tiling() found; leaving current state unchanged.") except Exception as e: print(f"⚠️ VAE tiling toggle failed: {e}") # ============================================================ # Derived conditioning (Depth only) — ViTPose REMOVED # ============================================================ DEPTH_MODEL_ID = "depth-anything/Depth-Anything-V2-Small-hf" _DEPTH_CACHE = {} def _derived_device(use_gpu: bool) -> torch.device: return torch.device("cuda" if (use_gpu and torch.cuda.is_available()) else "cpu") def _load_depth_models(dev: torch.device): key = str(dev) if key in _DEPTH_CACHE: return _DEPTH_CACHE[key] proc = AutoImageProcessor.from_pretrained(DEPTH_MODEL_ID) model = AutoModelForDepthEstimation.from_pretrained(DEPTH_MODEL_ID).to(dev) model.eval() _DEPTH_CACHE[key] = (proc, model) return _DEPTH_CACHE[key] def make_depth_map(img: Image.Image, *, use_gpu: bool) -> Image.Image: img = img.convert("RGB") dev = _derived_device(use_gpu) proc, model = _load_depth_models(dev) inputs = proc(images=img, return_tensors="pt") inputs = {k: v.to(dev) for k, v in inputs.items()} with torch.no_grad(): out = model(**inputs) pred = out.predicted_depth # (B,H,W) pred = torch.nn.functional.interpolate( pred.unsqueeze(1), size=(img.height, img.width), mode="bicubic", align_corners=False, ).squeeze(1)[0] arr = pred.detach().float().cpu().numpy() arr = arr - float(arr.min()) denom = float(arr.max()) + 1e-8 arr = arr / denom depth8 = (arr * 255.0).clip(0, 255).astype(np.uint8) return Image.fromarray(depth8, mode="L").convert("RGB") # ============================================================ # LoRA adapters + presets # ============================================================ NONE_LORA = "None" ADAPTER_SPECS = { "Consistance": { "type": "single", "repo": "Pr0f3ssi0n4ln00b/QIE_2511_Consistency_Lora", "weights": "qe2511_consis_alpha_patched.safetensors", "adapter_name": "Consistency", "strength": 0.6, }, "Semirealistic-photo-detailer": { "type": "single", "repo": "rzgar/Qwen-Image-Edit-semi-realistic-detailer", "weights": "Qwen-Image-Edit-Anime-Semi-Realistic-Detailer-v1.safetensors", "adapter_name": "semirealistic", "strength": 1.0, }, "AnyPose": { "type": "package", "requires_two_images": True, "image2_label": "Picture 2 (Pose Reference)", "parts": [ { "repo": "lilylilith/AnyPose", "weights": "2511-AnyPose-base-000006250.safetensors", "adapter_name": "anypose-base", "strength": 0.7, }, { "repo": "lilylilith/AnyPose", "weights": "2511-AnyPose-helper-00006000.safetensors", "adapter_name": "anypose-helper", "strength": 0.7, }, ], }, "Any2Real_2601": { "type": "single", "repo": "lrzjason/Anything2Real_2601", "weights": "anything2real_2601_A_final_patched.safetensors", "adapter_name": "photoreal", "strength": 1.0, }, "Hyperrealistic-Portrait": { "type": "single", "repo": "prithivMLmods/Qwen-Image-Edit-2511-Hyper-Realistic-Portrait", "weights": "HRP_20.safetensors", "adapter_name": "HRPortrait", "strength": 1.0, }, "Ultrarealistic-Portrait": { "type": "single", "repo": "prithivMLmods/Qwen-Image-Edit-2511-Ultra-Realistic-Portrait", "weights": "URP_20.safetensors", "adapter_name": "URPortrait", "strength": 1.0, }, "BFS-Best-FaceSwap": { "type": "single", "requires_two_images": True, "image2_label": "Picture 2 (Head/Face Donor)", "repo": "Alissonerdx/BFS-Best-Face-Swap", "weights": "bfs_head_v5_2511_original.safetensors", "adapter_name": "BFS-Best-Faceswap", "strength": 1.0, "needs_alpha_fix": True, }, "BFS-Best-FaceSwap-merge": { "type": "single", "requires_two_images": True, "image2_label": "Picture 2 (Head/Face Donor)", "repo": "Alissonerdx/BFS-Best-Face-Swap", "weights": "bfs_head_v5_2511_merged_version_rank_32_fp32.safetensors", "adapter_name": "BFS-Best-Faceswap-merge", "strength": 1.1, "needs_alpha_fix": True, }, "F2P": { "type": "single", "repo": "DiffSynth-Studio/Qwen-Image-Edit-F2P", "weights": "edit_0928_lora_step40000.safetensors", "adapter_name": "F2P", "strength": 1.0, }, "Multiple-Angles": { "type": "single", "repo": "dx8152/Qwen-Edit-2509-Multiple-angles", "weights": "镜头转换.safetensors", "adapter_name": "multiple-angles", "strength": 1.0, }, "Light-Restoration": { "type": "single", "repo": "dx8152/Qwen-Image-Edit-2509-Light_restoration", "weights": "移除光影.safetensors", "adapter_name": "light-restoration", "strength": 1.0, }, "Relight": { "type": "single", "repo": "dx8152/Qwen-Image-Edit-2509-Relight", "weights": "Qwen-Edit-Relight.safetensors", "adapter_name": "relight", "strength": 1.0, }, "Multi-Angle-Lighting": { "type": "single", "repo": "dx8152/Qwen-Edit-2509-Multi-Angle-Lighting", "weights": "多角度灯光-251116.safetensors", "adapter_name": "multi-angle-lighting", "strength": 1.0, }, "Edit-Skin": { "type": "single", "repo": "tlennon-ie/qwen-edit-skin", "weights": "qwen-edit-skin_1.1_000002750.safetensors", "adapter_name": "edit-skin", "strength": 1.0, }, "Next-Scene": { "type": "single", "repo": "lovis93/next-scene-qwen-image-lora-2509", "weights": "next-scene_lora-v2-3000.safetensors", "adapter_name": "next-scene", "strength": 1.0, }, "Flat-Log": { "type": "single", "repo": "tlennon-ie/QwenEdit2509-FlatLogColor", "weights": "QwenEdit2509-FlatLogColor.safetensors", "adapter_name": "flat-log", "strength": 1.0, }, "Upscale-Image": { "type": "single", "repo": "vafipas663/Qwen-Edit-2509-Upscale-LoRA", "weights": "qwen-edit-enhance_64-v3_000001000.safetensors", "adapter_name": "upscale-image", "strength": 1.0, }, "Upscale2K": { "type": "single", "repo": "valiantcat/Qwen-Image-Edit-2509-Upscale2K", "weights": "qwen_image_edit_2509_upscale.safetensors", "adapter_name": "upscale-2k", "strength": 1.0, "target_long_edge": 2048, }, } LORA_PRESET_PROMPTS = { "Any2Real_2601": "change the picture 1 to realistic photograph", "Semirealistic-photo-detailer": "transform the image to semi-realistic image", "AnyPose": ( "Make the person in image 1 do the exact same pose of the person in image 2. " "Changing the style and background of the image of the person in image 1 is undesirable, so don't do it. " "The new pose should be pixel accurate to the pose we are trying to copy. " "Change the field of view and angle to match exactly image 2." ), "Hyperrealistic-Portrait": ( "Transform the image into an ultra-realistic photorealistic portrait with strict identity preservation, " "facing straight to the camera. Enhance pore-level skin textures, realistic moisture effects, and natural wet hair clumping. " "Use shallow depth of field with a clean background." ), "Ultrarealistic-Portrait": ( "Transform the image into an ultra-realistic glamour portrait while strictly preserving the subject’s identity. " "Enhance cinematic directional lighting and keep realism without over-smoothing." ), "Upscale2K": "Upscale this picture to 4K resolution.", "BFS-Best-FaceSwap": ( "head_swap: start with Picture 1 as the base image. replace the head with Picture 2, preserving identity of Picture 2. " "copy eye direction and micro-expressions from Picture 1. high quality, sharp details, 4k" ), "BFS-Best-FaceSwap-merge": ( "head_swap: start with Picture 1 as the base image. replace the head with Picture 2, preserving identity of Picture 2. " "copy eye direction and micro-expressions from Picture 1. high quality, sharp details, 4k" ), } LOADED_ADAPTERS = set() # ============================================================ # Helpers: resolution # ============================================================ def _round_to_multiple(x: int, m: int) -> int: m = max(1, int(m)) return max(m, (int(x) // m) * m) def compute_canvas_dimensions_from_area(image: Image.Image, target_area: int, multiple_of: int) -> tuple[int, int]: w, h = image.size aspect = w / h if h else 1.0 from qwenimage.pipeline_qwenimage_edit_plus import calculate_dimensions width, height = calculate_dimensions(int(target_area), float(aspect), multiple=int(multiple_of)) width = _round_to_multiple(int(width), int(multiple_of)) height = _round_to_multiple(int(height), int(multiple_of)) return width, height def get_target_area_for_lora(image: Image.Image, lora_adapter: str, user_target_megapixels: float) -> int: spec = ADAPTER_SPECS.get(lora_adapter, {}) if "target_area" in spec: try: return int(spec["target_area"]) except Exception: pass if "target_megapixels" in spec: try: mp = float(spec["target_megapixels"]) return int(mp * 1024 * 1024) except Exception: pass if "target_long_edge" in spec: try: long_edge = int(spec["target_long_edge"]) w, h = image.size if w >= h: new_w = long_edge new_h = int(round(long_edge * (h / w))) else: new_h = long_edge new_w = int(round(long_edge * (w / h))) return int(new_w * new_h) except Exception: pass return int(float(user_target_megapixels) * 1024 * 1024) # ============================================================ # Helpers: gallery normalization # ============================================================ def _to_pil_rgb(x) -> Optional[Image.Image]: if x is None: return None if isinstance(x, tuple) and len(x) >= 1: x = x[0] if x is None: return None if isinstance(x, Image.Image): return x.convert("RGB") if isinstance(x, np.ndarray): return Image.fromarray(x).convert("RGB") try: return Image.fromarray(np.array(x)).convert("RGB") except Exception: return None def _append_to_gallery(existing, new_img: Image.Image): items = [] if existing: for it in existing: pil = _to_pil_rgb(it) if pil is not None: items.append(pil) items.append(new_img) return items def lora_requires_two_images(lora_adapter: str) -> bool: return bool(ADAPTER_SPECS.get(lora_adapter, {}).get("requires_two_images", False)) def image2_label_for_lora(lora_adapter: str) -> str: return str(ADAPTER_SPECS.get(lora_adapter, {}).get("image2_label", "Picture 2")) # ============================================================ # Helpers: BFS alpha key fix / strict filtering for merged safetensors # ============================================================ def _inject_missing_alpha_keys(state_dict: dict) -> dict: bases = {} for k, v in state_dict.items(): if not isinstance(v, torch.Tensor): continue if k.endswith(".lora_down.weight") and v.ndim >= 1: base = k[: -len(".lora_down.weight")] rank = int(v.shape[0]) bases[base] = rank for base, rank in bases.items(): alpha_tensor = torch.tensor(float(rank), dtype=torch.float32) full_alpha = f"{base}.alpha" if full_alpha not in state_dict: state_dict[full_alpha] = alpha_tensor if base.startswith("diffusion_model."): stripped_base = base[len("diffusion_model.") :] stripped_alpha = f"{stripped_base}.alpha" if stripped_alpha not in state_dict: state_dict[stripped_alpha] = alpha_tensor return state_dict def _filter_to_diffusers_lora_keys(state_dict: dict) -> tuple[dict, dict]: keep_suffixes = ( ".lora_up.weight", ".lora_down.weight", ".lora_mid.weight", ".alpha", ".lora_alpha", ) dropped_patch = 0 dropped_other = 0 kept = 0 normalized_alpha = 0 out = {} for k, v in state_dict.items(): if not isinstance(v, torch.Tensor): dropped_other += 1 continue if k.endswith(".diff") or k.endswith(".diff_b"): dropped_patch += 1 continue if not k.endswith(keep_suffixes): dropped_other += 1 continue if k.endswith(".lora_alpha"): base = k[: -len(".lora_alpha")] k2 = f"{base}.alpha" out[k2] = v.float() if v.dtype != torch.float32 else v normalized_alpha += 1 kept += 1 continue out[k] = v kept += 1 stats = { "kept": kept, "dropped_patch": dropped_patch, "dropped_other": dropped_other, "normalized_alpha": normalized_alpha, } return out, stats def _duplicate_stripped_prefix_keys(state_dict: dict, prefix: str = "diffusion_model.") -> dict: out = dict(state_dict) for k, v in list(state_dict.items()): if not k.startswith(prefix): continue stripped = k[len(prefix) :] if stripped not in out: out[stripped] = v return out def _load_lora_weights_with_fallback(repo: str, weight_name: str, adapter_name: str, needs_alpha_fix: bool = False): try: pipe.load_lora_weights(repo, weight_name=weight_name, adapter_name=adapter_name) return except (KeyError, ValueError) as e: if not needs_alpha_fix: raise print( "⚠️ LoRA load failed (will try safe dict fallback). " f"Adapter={adapter_name!r} file={weight_name!r} error={type(e).__name__}: {e}" ) local_path = hf_hub_download(repo_id=repo, filename=weight_name) sd = safetensors_load_file(local_path) sd = _inject_missing_alpha_keys(sd) sd, stats = _filter_to_diffusers_lora_keys(sd) sd = _duplicate_stripped_prefix_keys(sd) print("LoRA dict stats:", stats) pipe.load_lora_weights(sd, adapter_name=adapter_name) return def _ensure_loaded_and_get_active_adapters(selected_lora: str): spec = ADAPTER_SPECS.get(selected_lora) if not spec: raise gr.Error(f"Configuration not found for: {selected_lora}") adapter_names = [] adapter_weights = [] if spec.get("type") == "package": parts = spec.get("parts", []) if not parts: raise gr.Error(f"Package spec has no parts: {selected_lora}") for part in parts: repo = part["repo"] weights = part["weights"] name = part["adapter_name"] strength = float(part.get("strength", 1.0)) needs_alpha_fix = bool(part.get("needs_alpha_fix", False)) if name not in LOADED_ADAPTERS: _load_lora_weights_with_fallback(repo, weights, name, needs_alpha_fix=needs_alpha_fix) LOADED_ADAPTERS.add(name) adapter_names.append(name) adapter_weights.append(strength) else: repo = spec["repo"] weights = spec["weights"] name = spec["adapter_name"] strength = float(spec.get("strength", 1.0)) needs_alpha_fix = bool(spec.get("needs_alpha_fix", False)) if name not in LOADED_ADAPTERS: _load_lora_weights_with_fallback(repo, weights, name, needs_alpha_fix=needs_alpha_fix) LOADED_ADAPTERS.add(name) adapter_names.append(name) adapter_weights.append(strength) return adapter_names, adapter_weights # ============================================================ # UI helpers # ============================================================ def _fmt_img_info(img: Optional[Image.Image]) -> str: if img is None: return "—" w, h = img.size mp = (w * h) / (1024 * 1024) ar = (w / h) if h else 0 return f"**{w}×{h}** • **{mp:.2f} MP** • **AR {ar:.3f}**" def _bfs_tooltip(selected_lora: str) -> gr.Update: if selected_lora in ("BFS-Best-FaceSwap", "BFS-Best-FaceSwap-merge"): return gr.update( visible=True, value="ℹ️ **BFS FaceSwap:** Picture 1 = **Base** (scene), Picture 2 = **Donor** (head/face).", ) if selected_lora == "AnyPose": return gr.update( visible=True, value="ℹ️ **AnyPose:** Picture 1 = **Subject**, Picture 2 = **Pose reference**.", ) return gr.update(visible=False, value="") # ============================================================ # Inference # ============================================================ def _seed_everything(seed: int): random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed) @spaces.GPU def infer( img1: Image.Image, img2: Optional[Image.Image], extra_gallery, prompt: str, lora_adapter: str, seed: int, randomize_seed: bool, guidance_scale: float, steps: int, target_megapixels: float, use_input_area: bool, keep_2x_output: bool, vae_tiling: bool, extras_condition_only: bool, resolution_multiple: int, vae_ref_megapixels: float, use_depth: bool, derived_on_gpu: bool, ): if img1 is None: raise gr.Error("Picture 1 is required.") img1 = img1.convert("RGB") img2 = img2.convert("RGB") if img2 is not None else None # Seed if randomize_seed: seed = random.randint(0, MAX_SEED) seed = int(seed) % MAX_SEED _seed_everything(seed) # VAE tiling toggle _apply_vae_tiling(bool(vae_tiling)) # Load / activate LoRA if lora_adapter != NONE_LORA: adapter_names, adapter_weights = _ensure_loaded_and_get_active_adapters(lora_adapter) pipe.set_adapters(adapter_names, adapter_weights) else: try: pipe.set_adapters([]) except Exception: pass # Images list: Picture1, Picture2 (optional), extras..., derived (optional) images = [img1] base_count = 1 if lora_requires_two_images(lora_adapter): if img2 is None: raise gr.Error(f"{lora_adapter} requires Picture 2.") images.append(img2) base_count = 2 else: img2 = None # ignore if not needed extras = [] if extra_gallery: for it in extra_gallery: p = _to_pil_rgb(it) if p is not None: extras.append(p) images.extend(extras) derived_preview = None derived_index = None if use_depth: derived_preview = make_depth_map(img1, use_gpu=bool(derived_on_gpu)) images.append(derived_preview) derived_index = len(images) - 1 # Canvas sizing res_mult = int(resolution_multiple) if use_input_area or float(target_megapixels) <= 0.0: target_area = int(img1.width * img1.height) else: target_area = int(get_target_area_for_lora(img1, lora_adapter, float(target_megapixels))) base_w, base_h = compute_canvas_dimensions_from_area(img1, target_area, res_mult) # Generate at 2x, then downsample unless keep_2x_output gen_w, gen_h = int(base_w * 2), int(base_h * 2) # Extra refs routing (VAE vs conditioning-only) if extras_condition_only: vae_indices = list(range(base_count)) else: vae_indices = list(range(len(images))) # Derived depth should ALWAYS be conditioning-only if derived_index is not None and derived_index in vae_indices: vae_indices = [i for i in vae_indices if i != derived_index] # VAE ref size override for extras only vae_ref_area = None if float(vae_ref_megapixels) > 0.0: vae_ref_area = int(float(vae_ref_megapixels) * 1024 * 1024) # Run out = pipe( image=images, prompt=prompt, true_cfg_scale=float(guidance_scale), num_inference_steps=int(steps), width=int(gen_w), height=int(gen_h), pad_to_canvas=True, vae_image_indices=vae_indices, resolution_multiple=int(res_mult), vae_ref_area=vae_ref_area, vae_ref_start_index=int(base_count), generator=torch.Generator(device=device).manual_seed(seed), ) result = out.images[0] if hasattr(out, "images") else out[0][0] if isinstance(result, np.ndarray): result = Image.fromarray(result) result = result.convert("RGB") if not keep_2x_output: result = result.resize((base_w, base_h), Image.Resampling.LANCZOS) # Cleanup gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() return result, seed, derived_preview # ============================================================ # UI # ============================================================ def _on_lora_change(selected_lora: str): # Prompt preset preset = LORA_PRESET_PROMPTS.get(selected_lora, "") prompt_update = gr.update(value=preset) if preset else gr.update() # Picture 2 visibility/label if lora_requires_two_images(selected_lora): img2_update = gr.update(visible=True, label=image2_label_for_lora(selected_lora)) else: img2_update = gr.update(visible=True, label="Picture 2") # keep visible, but optional tooltip_update = _bfs_tooltip(selected_lora) return prompt_update, img2_update, tooltip_update def _out_to_pic1(out_img): return gr.update(value=out_img) def _out_to_pic2(out_img): return gr.update(value=out_img) def _out_to_extras(existing, out_img): if out_img is None: return gr.update() return gr.update(value=_append_to_gallery(existing, out_img)) with gr.Blocks(theme=orange_red_theme) as demo: gr.Markdown( f""" # Qwen Image Edit — Rapid AIO LoRAs (Merged) This experimental space for **QIE-2511** uses an extracted Rapid AIO transformer with LoRA support and extra routing features. **Enabled features** - Optional conditioning-only routing for extra reference latents - Uncapped canvas sizing (MP-based) + **2× generation with optional downsample** - Optional **VAE tiling** (for high resolutions) - Optional **Depth mapping** for conditioning - Optional output routing back to inputs **Active AIO version:** `{AIO_VERSION}` *(source: {AIO_VERSION_SOURCE})* """ ) with gr.Row(): with gr.Column(scale=1): img1 = gr.Image(label="Picture 1", type="pil") img1_info = gr.Markdown("—") img2 = gr.Image(label="Picture 2", type="pil") img2_info = gr.Markdown("—") bfs_tip = gr.Markdown(visible=False) extra_gallery = gr.Gallery( label="Extra references (optional)", columns=4, height=180, ) with gr.Row(): use_depth = gr.Checkbox(label="Use Depth conditioning (adds a derived reference)", value=False) derived_on_gpu = gr.Checkbox(label="Run depth on GPU (if available)", value=True) derived_preview = gr.Image(label="Derived conditioning preview", interactive=False, format="png") with gr.Column(scale=1): lora_adapter = gr.Dropdown( label="LoRA", choices=[NONE_LORA] + sorted(list(ADAPTER_SPECS.keys())), value=NONE_LORA, ) prompt = gr.Textbox(label="Prompt", lines=4, placeholder="Describe the edit…") with gr.Row(): steps = gr.Slider(1, 80, value=40, step=1, label="Steps") guidance = gr.Slider(1.0, 10.0, value=4.0, step=0.1, label="CFG (true_cfg_scale)") with gr.Row(): resolution_multiple = gr.Dropdown( label="Resolution step (LCD lattice)", choices=[32, 56, 112], value=32, ) vae_ref_megapixels = gr.Slider( 0.0, 4.0, value=0.0, step=0.1, label="VAE ref MP override (extras only, 0 = off)" ) with gr.Row(): target_megapixels = gr.Slider( 0.0, 12.0, value=1.0, step=0.1, label="Canvas megapixels (0 = same as Picture 1)" ) use_input_area = gr.Checkbox(label="Use Picture 1 pixel area", value=False) with gr.Row(): keep_2x_output = gr.Checkbox(label="Keep 2× output (otherwise downsample)", value=False) extras_condition_only = gr.Checkbox(label="Route extras as conditioning-only (no VAE)", value=True) with gr.Row(): vae_tiling = gr.Checkbox(label="VAE tiling", value=False) randomize_seed = gr.Checkbox(label="Randomize seed", value=True) seed = gr.Number(label="Seed", value=0, precision=0) run_btn = gr.Button("Run", variant="primary") out_img = gr.Image(label="Output", type="pil") with gr.Row(): to_pic1 = gr.Button("Output → Picture 1") to_pic2 = gr.Button("Output → Picture 2") to_extras = gr.Button("Output → Extras (append)") # Live info updates img1.change(lambda x: _fmt_img_info(x), inputs=[img1], outputs=[img1_info]) img2.change(lambda x: _fmt_img_info(x), inputs=[img2], outputs=[img2_info]) # LoRA change lora_adapter.change(_on_lora_change, inputs=[lora_adapter], outputs=[prompt, img2, bfs_tip]) # Run run_btn.click( infer, inputs=[ img1, img2, extra_gallery, prompt, lora_adapter, seed, randomize_seed, guidance, steps, target_megapixels, use_input_area, keep_2x_output, vae_tiling, extras_condition_only, resolution_multiple, vae_ref_megapixels, use_depth, derived_on_gpu, ], outputs=[out_img, seed, derived_preview], ) # Output routing buttons to_pic1.click(_out_to_pic1, inputs=[out_img], outputs=[img1]) to_pic2.click(_out_to_pic2, inputs=[out_img], outputs=[img2]) to_extras.click(_out_to_extras, inputs=[extra_gallery, out_img], outputs=[extra_gallery]) demo.queue(max_size=32).launch()