Update app.py

app.py

@@ -2,16 +2,19 @@ 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
+import random
+from PIL import Image, ImageDraw
+from typing import Iterable, Optional
+
+from transformers import (
+    AutoImageProcessor,
+    AutoModelForDepthEstimation,
+)
 
-from transformers import AutoImageProcessor, AutoModelForDepthEstimation
 from huggingface_hub import hf_hub_download
 from safetensors.torch import load_file as safetensors_load_file
 
@@ -37,6 +40,7 @@ colors.orange_red = colors.Color(
     c950="#802200",
 )
 
+
 class OrangeRedTheme(Soft):
     def __init__(
         self,
@@ -92,6 +96,7 @@ class OrangeRedTheme(Soft):
             block_label_background_fill="*primary_200",
         )
 
+
 orange_red_theme = OrangeRedTheme()
 
 # ============================================================
@@ -99,6 +104,7 @@ orange_red_theme = OrangeRedTheme()
 # ============================================================
 
 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)
@@ -109,18 +115,17 @@ if torch.cuda.is_available():
     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
@@ -129,14 +134,18 @@ def _normalize_version(raw: str) -> Optional[str]:
         return None
     if _VER_RE.fullmatch(s):
         return s
+    # forgiving: allow "21" -> "v21"
     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})")
@@ -150,9 +159,12 @@ from qwenimage.pipeline_qwenimage_edit_plus import QwenImageEditPlusPipeline
 from qwenimage.transformer_qwenimage import QwenImageTransformer2DModel
 from qwenimage.qwen_fa3_processor import QwenDoubleStreamAttnProcessorFA3
 
+dtype = torch.bfloat16
+
+
 def _load_pipe_with_version(version: str) -> QwenImageEditPlusPipeline:
     sub = f"{version}/transformer"
-    print(f"Loading AIO transformer: {AIO_REPO_ID} / {sub}")
+    print(f"📦 Loading AIO transformer: {AIO_REPO_ID} / {sub}")
     p = QwenImageEditPlusPipeline.from_pretrained(
         "Qwen/Qwen-Image-Edit-2511",
         transformer=QwenImageTransformer2DModel.from_pretrained(
@@ -165,9 +177,11 @@ def _load_pipe_with_version(version: str) -> QwenImageEditPlusPipeline:
     ).to(device)
     return p
 
+
+# Forgiving load: try env/default version, fallback to v19 if it fails
 try:
     pipe = _load_pipe_with_version(AIO_VERSION)
-except Exception:
+except Exception as e:
     print("❌ Failed to load requested AIO_VERSION. Falling back to v19.")
     print("---- exception ----")
     print(traceback.format_exc())
@@ -176,12 +190,15 @@ except Exception:
     AIO_VERSION_SOURCE = "fallback_to_v19"
     pipe = _load_pipe_with_version(AIO_VERSION)
 
+# Apply FA3 Optimization
 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}")
 
+MAX_SEED = np.iinfo(np.int32).max
+
 # ============================================================
 # VAE tiling toggle (UI-controlled; OFF by default)
 # ============================================================
@@ -189,7 +206,9 @@ except Exception as e:
 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.
+    Note: this is global process state, so concurrent users could flip it between runs.
     """
     try:
         if enabled:
@@ -204,20 +223,25 @@ def _apply_vae_tiling(enabled: bool):
         else:
             if hasattr(pipe, "disable_vae_tiling"):
                 pipe.disable_vae_tiling()
-                print("VAE tiling DISABLED (per UI).")
+                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).")
+                print("🛑 VAE tiling DISABLED via pipe.vae.disable_tiling() (per UI).")
             else:
+                # If no disable method exists, we leave current state unchanged.
                 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
+# Derived conditioning (Transformers): Depth
 # ============================================================
+# Depth uses Depth Anything V2 Small (Transformers-compatible):
+# https://huggingface.co/depth-anything/Depth-Anything-V2-Small-hf
 
 DEPTH_MODEL_ID = "depth-anything/Depth-Anything-V2-Small-hf"
+
+# Lazy cache keyed by device string ("cpu" / "cuda")
 _DEPTH_CACHE = {}
 
 def _derived_device(use_gpu: bool) -> torch.device:
@@ -233,32 +257,27 @@ def _load_depth_models(dev: torch.device):
     _DEPTH_CACHE[key] = (proc, model)
     return _DEPTH_CACHE[key]
 
+@torch.inference_mode()
 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)
+    w, h = img.size
+    inputs = proc(images=img.convert("RGB"), return_tensors="pt").to(dev)
+    outputs = model(**inputs)
+    predicted = outputs.predicted_depth  # [B, H, W]
 
-    pred = out.predicted_depth  # (B,H,W)
-    pred = torch.nn.functional.interpolate(
-        pred.unsqueeze(1),
-        size=(img.height, img.width),
+    depth = torch.nn.functional.interpolate(
+        predicted.unsqueeze(1),
+        size=(h, w),
         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")
+    depth = depth - depth.min()
+    depth = depth / (depth.max() + 1e-8)
+    depth = (depth * 255.0).clamp(0, 255).to(torch.uint8).cpu().numpy()
+    return Image.fromarray(depth).convert("RGB")
 
 # ============================================================
 # LoRA adapters + presets
@@ -284,7 +303,7 @@ ADAPTER_SPECS = {
     "AnyPose": {
         "type": "package",
         "requires_two_images": True,
-        "image2_label": "Picture 2 (Pose Reference)",
+        "image2_label": "Upload Pose Reference (Image 2)",
         "parts": [
             {
                 "repo": "lilylilith/AnyPose",
@@ -324,22 +343,22 @@ ADAPTER_SPECS = {
     "BFS-Best-FaceSwap": {
         "type": "single",
         "requires_two_images": True,
-        "image2_label": "Picture 2 (Head/Face Donor)",
+        "image2_label": "Upload Head/Face Donor (Image 2)",
         "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,
+        "needs_alpha_fix": True,  # <-- fixes KeyError 'img_in.alpha'
     },
     "BFS-Best-FaceSwap-merge": {
         "type": "single",
         "requires_two_images": True,
-        "image2_label": "Picture 2 (Head/Face Donor)",
+        "image2_label": "Upload Head/Face Donor (Image 2)",
         "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,
+        "needs_alpha_fix": True,  # <-- fixes KeyError 'img_in.alpha'
     },
     "F2P": {
         "type": "single",
@@ -417,64 +436,75 @@ ADAPTER_SPECS = {
 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."
-    ),
+    "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. The position of the arms and head and legs should be the same as the pose we are trying to copy. Change the field of view and angle to match exactly image 2. Head tilt and eye gaze pose should match the person in 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 against the skin. Apply cool-toned soft-box lighting with subtle highlights and shadows, maintain realistic green-hazel eye catchlights without synthetic gloss, and preserve soft natural lip texture. Use shallow depth of field with a clean bokeh background, an 85mm macro photographic look, and raw photo grading without retouching to maintain realism and original details.",
+    "Ultrarealistic-Portrait": "Transform the image into an ultra-realistic glamour portrait while strictly preserving the subject’s identity. Apply a close-up composition with a slight head tilt and a hand near the face, enhance cinematic directional lighting with dramatic fashion-style highlights, and refine makeup details including glowing skin, glossy lips, luminous highlighter, and defined eyes. Increase skin realism with detailed epidermal textures such as micropores, microhairs, subtle oil sheen, natural highlights, soft wrinkles, and subsurface scattering. Maintain a luxury fashion-magazine look in a 9:16 aspect ratio, preserving realism, facial structure, and original details without over-smoothing or retouching.",
     "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"
-    ),
+    "BFS-Best-FaceSwap": "head_swap: start with Picture 1 as the base image, keeping its lighting, environment, and background. remove the head from Picture 1 completely and replace it with the head from Picture 2, strictly preserving the hair, eye color, and nose structure of Picture 2. copy the eye direction, head rotation, 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, keeping its lighting, environment, and background. remove the head from Picture 1 completely and replace it with the head from Picture 2, strictly preserving the hair, eye color, and nose structure of Picture 2. copy the eye direction, head rotation, and micro-expressions from Picture 1. high quality, sharp details, 4k",
 }
 
+# Track what is currently loaded in memory (adapter_name values)
 LOADED_ADAPTERS = set()
 
 # ============================================================
 # Helpers: resolution
 # ============================================================
 
+# We prefer *area-based* sizing (≈ megapixels) over long-edge sizing.
+# This aligns better with Qwen-Image-Edit's internal assumptions and reduces FOV drift.
+
 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]:
+def compute_canvas_dimensions_from_area(
+    image: Image.Image,
+    target_area: int,
+    multiple_of: int,
+) -> tuple[int, int]:
+    """Compute (width, height) that matches image aspect ratio and approximates target_area.
+
+    The result is floored to be divisible by multiple_of (typically vae_scale_factor*2).
+    """
     w, h = image.size
     aspect = w / h if h else 1.0
+
+    # Use the pipeline's own area->(w,h) helper for consistency.
     from qwenimage.pipeline_qwenimage_edit_plus import calculate_dimensions
-    width, height = calculate_dimensions(int(target_area), float(aspect), multiple=int(multiple_of))
+
+    width, height = calculate_dimensions(int(target_area), float(aspect))
     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:
+def get_target_area_for_lora(
+    image: Image.Image,
+    lora_adapter: str,
+    user_target_megapixels: float,
+) -> int:
+    """Return target pixel area for the canvas.
+
+    Priority:
+      1) Adapter spec: target_area (pixels) or target_megapixels
+      2) Adapter spec: target_long_edge (legacy) -> converted to area using image aspect
+      3) User slider target megapixels
+    """
     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
+
+    # Legacy support (e.g. Upscale2K)
     if "target_long_edge" in spec:
         try:
             long_edge = int(spec["target_long_edge"])
@@ -488,50 +518,109 @@ def get_target_area_for_lora(image: Image.Image, lora_adapter: str, user_target_
             return int(new_w * new_h)
         except Exception:
             pass
-    return int(float(user_target_megapixels) * 1024 * 1024)
+
+    # User default
+    try:
+        mp = float(user_target_megapixels)
+    except Exception:
+        mp = 1.0
+
+    # Treat 0 MP as "match input area"
+    if mp <= 0:
+        w, h = image.size
+        return int(w * h)
+
+    return int(mp * 1024 * 1024)
 
 # ============================================================
-# Helpers: gallery normalization
+# Helpers: multi-input routing + gallery normalization
 # ============================================================
 
+
+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", "Upload Reference (Image 2)"))
+
+
 def _to_pil_rgb(x) -> Optional[Image.Image]:
+    """
+    Accepts PIL / numpy / (image, caption) tuples from gr.Gallery and returns PIL RGB.
+    Gradio Gallery commonly yields tuples like (image, caption).
+    """
     if x is None:
         return None
+
+    # Gallery often returns (image, caption)
     if isinstance(x, tuple) and len(x) >= 1:
         x = x[0]
-    if x is None:
-        return None
+        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")
+
+    # Best-effort fallback
     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 build_labeled_images(
+    img1: Image.Image,
+    img2: Optional[Image.Image],
+    extra_imgs: Optional[list[Image.Image]],
+) -> dict[str, Image.Image]:
+    """
+    Creates labels image_1, image_2, image_3... based on what is actually uploaded:
+      - img1 is always image_1
+      - img2 becomes image_2 only if present
+      - extras start immediately after the last present base box
+    The pipeline receives images in this exact order.
+    """
+    labeled: dict[str, Image.Image] = {}
+    idx = 1
+
+    labeled[f"image_{idx}"] = img1
+    idx += 1
+
+    if img2 is not None:
+        labeled[f"image_{idx}"] = img2
+        idx += 1
+
+    if extra_imgs:
+        for im in extra_imgs:
+            if im is None:
+                continue
+            labeled[f"image_{idx}"] = im
+            idx += 1
+
+    return labeled
 
-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
+# Helpers: BFS alpha key fix
 # ============================================================
 
+
 def _inject_missing_alpha_keys(state_dict: dict) -> dict:
+    """
+    Diffusers' Qwen LoRA converter expects '<module>.alpha' keys.
+    BFS safetensors omits them. We inject alpha = rank (neutral scaling).
+
+    IMPORTANT: diffusers may strip 'diffusion_model.' before lookup, so we
+    inject BOTH:
+      - diffusion_model.xxx.alpha
+      - xxx.alpha
+    """
     bases = {}
+
     for k, v in state_dict.items():
         if not isinstance(v, torch.Tensor):
             continue
@@ -542,17 +631,37 @@ def _inject_missing_alpha_keys(state_dict: dict) -> dict:
 
     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]:
+    """Return (filtered_state_dict, stats).
+
+    Some ComfyUI/Qwen safetensors (especially "merged" variants) include non-LoRA
+    delta/patch keys like `*.diff` and `*.diff_b` alongside real LoRA tensors.
+    Diffusers' internal Qwen LoRA converter is strict: any leftover keys cause an
+    error (`state_dict should be empty...`).
+
+    This helper keeps only the keys Diffusers can consume as a LoRA:
+      - `*.lora_up.weight`
+      - `*.lora_down.weight`
+      - (rare) `*.lora_mid.weight`
+      - alpha keys: `*.alpha` (or `*.lora_alpha` which we normalize to `*.alpha`)
+
+    It also drops known patch keys (`*.diff`, `*.diff_b`) and everything else.
+    """
+
     keep_suffixes = (
         ".lora_up.weight",
         ".lora_down.weight",
@@ -560,29 +669,37 @@ def _filter_to_diffusers_lora_keys(state_dict: dict) -> tuple[dict, dict]:
         ".alpha",
         ".lora_alpha",
     )
+
     dropped_patch = 0
     dropped_other = 0
     kept = 0
     normalized_alpha = 0
 
-    out = {}
+    out: dict[str, torch.Tensor] = {}
     for k, v in state_dict.items():
         if not isinstance(v, torch.Tensor):
+            # Ignore non-tensor entries if any.
             dropped_other += 1
             continue
+
+        # Drop ComfyUI "delta" keys that Diffusers' LoRA loader will never consume.
         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"):
+            # Normalize common alt name to what Diffusers expects.
             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
 
@@ -594,7 +711,15 @@ def _filter_to_diffusers_lora_keys(state_dict: dict) -> tuple[dict, dict]:
     }
     return out, stats
 
+
 def _duplicate_stripped_prefix_keys(state_dict: dict, prefix: str = "diffusion_model.") -> dict:
+    """Ensure both prefixed and unprefixed variants exist for LoRA-related keys.
+
+    Diffusers' Qwen LoRA conversion may strip `diffusion_model.` when looking up
+    modules. Some exports only include prefixed keys. To be maximally compatible,
+    we duplicate LoRA keys (and alpha) in stripped form when missing.
+    """
+
     out = dict(state_dict)
     for k, v in list(state_dict.items()):
         if not k.startswith(prefix):
@@ -604,11 +729,18 @@ def _duplicate_stripped_prefix_keys(state_dict: dict, prefix: str = "diffusion_m
             out[stripped] = v
     return out
 
+
 def _load_lora_weights_with_fallback(repo: str, weight_name: str, adapter_name: str, needs_alpha_fix: bool = False):
+    """
+    Normal path: pipe.load_lora_weights(repo, weight_name=..., adapter_name=...)
+    BFS fallback: download safetensors, inject missing alpha keys, then load from dict.
+    """
     try:
         pipe.load_lora_weights(repo, weight_name=weight_name, adapter_name=adapter_name)
         return
     except (KeyError, ValueError) as e:
+        # KeyError: missing required alpha keys (common in BFS)
+        # ValueError: Diffusers Qwen converter found leftover keys (e.g. .diff/.diff_b)
         if not needs_alpha_fix:
             raise
 
@@ -616,15 +748,34 @@ def _load_lora_weights_with_fallback(repo: str, weight_name: str, adapter_name:
             "⚠️ 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)
+
+        # 1) Inject required `<module>.alpha` keys (neutral scaling alpha=rank).
         sd = _inject_missing_alpha_keys(sd)
+
+        # 2) Keep only LoRA + alpha keys; drop ComfyUI patch/delta keys.
         sd, stats = _filter_to_diffusers_lora_keys(sd)
+
+        # 3) Duplicate stripped keys (remove `diffusion_model.`) for compatibility.
         sd = _duplicate_stripped_prefix_keys(sd)
-        print("LoRA dict stats:", stats)
+
+        print(
+            "🧹 LoRA dict cleanup stats: "
+            f"kept={stats['kept']} dropped_patch={stats['dropped_patch']} "
+            f"dropped_other={stats['dropped_other']} normalized_alpha={stats['normalized_alpha']}"
+        )
+
         pipe.load_lora_weights(sd, adapter_name=adapter_name)
         return
 
+
+# ============================================================
+# LoRA loader: single/package + strengths
+# ============================================================
+
+
 def _ensure_loaded_and_get_active_adapters(selected_lora: str):
     spec = ADAPTER_SPECS.get(selected_lora)
     if not spec:
@@ -637,356 +788,522 @@ def _ensure_loaded_and_get_active_adapters(selected_lora: str):
         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"]
+            adapter_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)
+            if adapter_name not in LOADED_ADAPTERS:
+                print(f"--- Downloading and Loading Adapter Part: {selected_lora} / {adapter_name} ---")
+                try:
+                    _load_lora_weights_with_fallback(
+                        repo=repo,
+                        weight_name=weights,
+                        adapter_name=adapter_name,
+                        needs_alpha_fix=needs_alpha_fix,
+                    )
+                    LOADED_ADAPTERS.add(adapter_name)
+                except Exception as e:
+                    raise gr.Error(f"Failed to load adapter part {selected_lora}/{adapter_name}: {e}")
+            else:
+                print(f"--- Adapter part already loaded: {selected_lora} / {adapter_name} ---")
 
-            adapter_names.append(name)
+            adapter_names.append(adapter_name)
             adapter_weights.append(strength)
 
     else:
         repo = spec["repo"]
         weights = spec["weights"]
-        name = spec["adapter_name"]
+        adapter_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)
+        if adapter_name not in LOADED_ADAPTERS:
+            print(f"--- Downloading and Loading Adapter: {selected_lora} ---")
+            try:
+                _load_lora_weights_with_fallback(
+                    repo=repo,
+                    weight_name=weights,
+                    adapter_name=adapter_name,
+                    needs_alpha_fix=needs_alpha_fix,
+                )
+                LOADED_ADAPTERS.add(adapter_name)
+            except Exception as e:
+                raise gr.Error(f"Failed to load adapter {selected_lora}: {e}")
+        else:
+            print(f"--- Adapter {selected_lora} is already loaded. ---")
 
-        adapter_names.append(name)
-        adapter_weights.append(strength)
+        adapter_names = [adapter_name]
+        adapter_weights = [strength]
 
     return adapter_names, adapter_weights
 
+
 # ============================================================
-# UI helpers
+# UI handlers
 # ============================================================
 
-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="")
 
+def on_lora_change_ui(selected_lora, current_prompt, current_extras_condition_only):
+    # Preset prompt (fill only if empty)
+    if selected_lora != NONE_LORA:
+        preset = LORA_PRESET_PROMPTS.get(selected_lora, "")
+        if preset and (current_prompt is None or str(current_prompt).strip() == ""):
+            prompt_update = gr.update(value=preset)
+        else:
+            prompt_update = gr.update(value=current_prompt)
+    else:
+        prompt_update = gr.update(value=current_prompt)
+
+    # Image2 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=False, value=None, label="Upload Reference (Image 2)")
+
+    # Extra references routing default:
+    # For BFS/AnyPose-like adapters, it's usually safer to keep extra refs as conditioning-only.
+    if selected_lora in ("BFS-Best-FaceSwap", "BFS-Best-FaceSwap-merge", "AnyPose"):
+        extras_update = gr.update(value=True)
+    else:
+        extras_update = gr.update(value=current_extras_condition_only)
+
+    return prompt_update, img2_update, extras_update
 # ============================================================
-# Inference
+# UI helpers: output routing + derived conditioning
 # ============================================================
 
-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)
+def set_output_as_image1(last):
+    if last is None:
+        raise gr.Error("No output available yet.")
+    return gr.update(value=last)
 
-@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
+def set_output_as_image2(last):
+    if last is None:
+        raise gr.Error("No output available yet.")
+    return gr.update(value=last)
 
-    # 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))
+def set_output_as_extra(last, existing_extra):
+    if last is None:
+        raise gr.Error("No output available yet.")
+    return _append_to_gallery(existing_extra, last)
 
-    # 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)))
+@spaces.GPU
+def add_derived_ref(img1, existing_extra, derived_type, derived_use_gpu):
+    if img1 is None:
+        raise gr.Error("Please upload Image 1 first.")
 
-    base_w, base_h = compute_canvas_dimensions_from_area(img1, target_area, res_mult)
+    if derived_type == "None":
+        return gr.update(value=existing_extra), gr.update(visible=False, value=None)
 
-    # Generate at 2x, then downsample unless keep_2x_output
-    gen_w, gen_h = int(base_w * 2), int(base_h * 2)
+    base = img1.convert("RGB")
 
-    # Extra refs routing (VAE vs conditioning-only)
-    if extras_condition_only:
-        vae_indices = list(range(base_count))
+    if derived_type == "Depth (Depth Anything V2 Small)":
+        derived = make_depth_map(base, use_gpu=bool(derived_use_gpu))
     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),
-    )
+        raise gr.Error(f"Unknown derived type: {derived_type}")
 
-    result = out.images[0] if hasattr(out, "images") else out[0][0]
-    if isinstance(result, np.ndarray):
-        result = Image.fromarray(result)
+    new_gallery = _append_to_gallery(existing_extra, derived)
+    return gr.update(value=new_gallery), gr.update(visible=True, value=derived)
 
-    result = result.convert("RGB")
 
-    if not keep_2x_output:
-        result = result.resize((base_w, base_h), Image.Resampling.LANCZOS)
+# ============================================================
+# Inference
+# ============================================================
 
-    # Cleanup
+
+@spaces.GPU
+def infer(
+    input_image_1,
+    input_image_2,
+    input_images_extra,  # gallery multi-image box
+    prompt,
+    lora_adapter,
+    seed,
+    randomize_seed,
+    guidance_scale,
+    steps,
+    target_megapixels,
+    extras_condition_only,
+    pad_to_canvas,
+    vae_tiling,  # VAE tiling toggle
+    resolution_multiple,
+    vae_ref_megapixels,
+    decoder_vae,
+    keep_decoder_2x,
+    progress=gr.Progress(track_tqdm=True),
+):
     gc.collect()
     if torch.cuda.is_available():
         torch.cuda.empty_cache()
 
-    return result, seed, derived_preview
+    if input_image_1 is None:
+        raise gr.Error("Please upload Image 1.")
 
-# ============================================================
-# UI
-# ============================================================
+    # Handle "None"
+    if lora_adapter == NONE_LORA:
+        try:
+            pipe.set_adapters([], adapter_weights=[])
+        except Exception:
+            if LOADED_ADAPTERS:
+                pipe.set_adapters(list(LOADED_ADAPTERS), adapter_weights=[0.0] * len(LOADED_ADAPTERS))
+    else:
+        adapter_names, adapter_weights = _ensure_loaded_and_get_active_adapters(lora_adapter)
+        pipe.set_adapters(adapter_names, adapter_weights=adapter_weights)
 
-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()
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
 
-    # 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)
+    generator = torch.Generator(device=device).manual_seed(seed)
+    negative_prompt = (
+        "worst quality, low quality, bad anatomy, bad hands, text, error, missing fingers, "
+        "extra digit, fewer digits, cropped, jpeg artifacts, signature, watermark, username, blurry"
+    )
+
+    img1 = input_image_1.convert("RGB")
+    img2 = input_image_2.convert("RGB") if input_image_2 is not None else None
+
+    # Normalize extra images (Gallery) to PIL RGB (handles tuples from Gallery)
+    extra_imgs: list[Image.Image] = []
+    if input_images_extra:
+        for item in input_images_extra:
+            pil = _to_pil_rgb(item)
+            if pil is not None:
+                extra_imgs.append(pil)
+
+    # Enforce existing 2-image LoRA behavior (image_1 + image_2 required)
+    if lora_requires_two_images(lora_adapter) and img2 is None:
+        raise gr.Error("This LoRA needs two images. Please upload Image 2 as well.")
+
+    # Label images as image_1, image_2, image_3...
+    labeled = build_labeled_images(img1, img2, extra_imgs)
+
+    # Pass to pipeline in labeled order. Keep single-image call when only one is present.
+    pipe_images = list(labeled.values())
+    if len(pipe_images) == 1:
+        pipe_images = pipe_images[0]
+
+    # Resolution derived from Image 1 (base/body/target)
+    # Use target *area* (≈ megapixels) rather than long-edge sizing to reduce FOV drift.
+    target_area = get_target_area_for_lora(img1, lora_adapter, float(target_megapixels))
+    width, height = compute_canvas_dimensions_from_area(
+        img1,
+        target_area=target_area,
+        multiple_of=int(resolution_multiple),
+    )
+
+    # Decide which images participate in the VAE latent stream.
+    # If enabled, extra references beyond (Img_1, Img_2) become conditioning-only.
+    vae_image_indices = None
+    if extras_condition_only:
+        if isinstance(pipe_images, list) and len(pipe_images) > 2:
+            vae_image_indices = [0, 1] if len(pipe_images) >= 2 else [0]
 
-    return prompt_update, img2_update, tooltip_update
+    try:
+        print(
+             "[DEBUG][infer] submitting request | "
+             f"lora_adapter={lora_adapter!r} seed={seed} prompt={prompt!r}"
+        )
+        print(f"[DEBUG][infer] canvas={width}x{height} (~{(width*height)/1_048_576:.3f} MP) vae_tiling={bool(vae_tiling)}")
 
-def _out_to_pic1(out_img):
-    return gr.update(value=out_img)
+        # ✅ Apply UI toggle per-request (OFF by default)
+        # Lattice multiple passed to pipeline too (anti-drift / valid size grid)
+        res_mult = int(resolution_multiple) if resolution_multiple is not None else int(pipe.vae_scale_factor * 2)
 
-def _out_to_pic2(out_img):
-    return gr.update(value=out_img)
+        # Optional: override VAE sizing for *extra* references (beyond Image 1 / Image 2)
+        # Interpreted as megapixels; 0 disables override (uses canvas).
+        try:
+            mp_ref = float(vae_ref_megapixels)
+        except Exception:
+            mp_ref = 0.0
+
+        vae_ref_area = int(mp_ref * 1024 * 1024) if mp_ref and mp_ref > 0 else None
+
+        # Extras start index depends on whether Image 2 exists
+        base_ref_count = 2 if img2 is not None else 1
+
+        _apply_vae_tiling(bool(vae_tiling))
+
+        result = pipe(
+            image=pipe_images,
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            height=height,
+            width=width,
+            num_inference_steps=steps,
+            generator=generator,
+            true_cfg_scale=guidance_scale,
+            vae_image_indices=vae_image_indices,
+            pad_to_canvas=bool(pad_to_canvas),
+            resolution_multiple=res_mult,
+            vae_ref_area=vae_ref_area,
+            vae_ref_start_index=base_ref_count,
+            decoder_vae=str(decoder_vae).lower(),
+            keep_decoder_2x=bool(keep_decoder_2x),
+        ).images[0]
+        return result, seed, result
+    finally:
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
 
-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.
+@spaces.GPU
+def infer_example(input_image, prompt, lora_adapter):
+    if input_image is None:
+        return None, 0, None
+    input_pil = input_image.convert("RGB")
+    guidance_scale = 1.0
+    steps = 4
+    # Examples don't supply Image 2 or extra images; and example list doesn't include AnyPose/BFS.
+    # Keep VAE tiling OFF in examples (matches default).
+    result, seed, last = infer(
+        input_pil,
+        None,
+        None,
+        prompt,
+        lora_adapter,
+        0,
+        True,
+        guidance_scale,
+        steps,
+        1.0,
+        True,
+        True,
+        False,  # vae_tiling
+    )
+    return result, seed, last
 
-**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})*
+# ============================================================
+# UI
+# ============================================================
+
+css = """
+#col-container {
+    margin: 0 auto;
+    max-width: 960px;
+}
+#main-title h1 {font-size: 2.1em !important;}
 """
-    )
 
-    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)"
+aio_status_line = (
+    f"**AIO transformer version:** `{AIO_VERSION}`  "
+    f"({AIO_VERSION_SOURCE}; env `AIO_VERSION`={_AIO_ENV_RAW!r})"
+)
+
+with gr.Blocks() as demo:
+    with gr.Column(elem_id="col-container"):
+        gr.Markdown("# **Qwen-Image-Edit-2511-LoRAs-Fast**", elem_id="main-title")
+        gr.Markdown(
+            f"""This experimental space for QIE-2511 utilizes extracted transformers of Phr00t’s Rapid AIO merge with LoRA support and a couple of extra features:
+
+- Optional conditioning-only routing for extra reference latents
+- Uncapped canvas resolution
+- Optional VAE tiling for high resolutions
+- Optional depth mapping for conditioning
+- Optional routing of output to input for further iterations
+- Optional alternative decoder VAE (Wan2x 2×)
+
+Current environment is running **{AIO_VERSION}** of the Rapid AIO. Duplicate the space and set the **AIO_VERSION** space variable to use a different version."""
+        )
+        gr.Markdown(aio_status_line)
+
+        with gr.Row(equal_height=True):
+            with gr.Column():
+                input_image_1 = gr.Image(label="Upload Image 1 (Base / Target)", type="pil", height=290)
+                input_image_2 = gr.Image(label="Upload Reference (Image 2)", type="pil", height=290, visible=False)
+
+                input_images_extra = gr.Gallery(
+                    label="Upload Additional Images (auto-indexed after Image 1/2)",
+                    type="pil",
+                    height=290,
+                    columns=4,
+                    rows=2,
+                    interactive=True,
                 )
 
-            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)"
+                prompt = gr.Text(
+                    label="Edit Prompt",
+                    show_label=True,
+                    placeholder="e.g., transform into photo..",
                 )
-                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)
+                run_button = gr.Button("Edit Image", variant="primary")
 
-            with gr.Row():
-                vae_tiling = gr.Checkbox(label="VAE tiling", value=False)
-                randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
+            with gr.Column():
+                output_image = gr.Image(label="Output Image", interactive=False, format="png", height=353)
 
-            seed = gr.Number(label="Seed", value=0, precision=0)
+                last_output = gr.State(value=None)
 
-            run_btn = gr.Button("Run", variant="primary")
-            out_img = gr.Image(label="Output", type="pil")
+                with gr.Row():
+                    btn_out_to_img1 = gr.Button("⬅️ Output → Image 1", variant="secondary")
+                    btn_out_to_img2 = gr.Button("⬅️ Output → Image 2", variant="secondary")
+                    btn_out_to_extra = gr.Button("➕ Output → Extra Ref", variant="secondary")
 
-            with gr.Row():
-                to_pic1 = gr.Button("Output → Picture 1")
-                to_pic2 = gr.Button("Output → Picture 2")
-                to_extras = gr.Button("Output → Extras (append)")
+                derived_preview = gr.Image(
+                    label="Derived Conditioning Preview",
+                    interactive=False,
+                    format="png",
+                    height=200,
+                    visible=False,
+                )
 
-    # 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])
+                with gr.Row():
+                    lora_choices = [NONE_LORA] + list(ADAPTER_SPECS.keys())
+                    lora_adapter = gr.Dropdown(
+                        label="Choose Editing Style",
+                        choices=lora_choices,
+                        value=NONE_LORA,
+                    )
+
+                with gr.Accordion("Advanced Settings", open=False, visible=True):
+                    with gr.Accordion("Derived Conditioning (Pose / Depth)", open=False):
+                        derived_type = gr.Dropdown(
+                            label="Derived Type (from Image 1)",
+                            choices=["None", "Depth (Depth Anything V2 Small)"],
+                            value="None",
+                        )
+                        derived_use_gpu = gr.Checkbox(label="Use GPU for derived model", value=False)
+                        add_derived_btn = gr.Button("➕ Add derived ref to Extras (conditioning-only recommended)")
+
+                    seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0)
+                    randomize_seed = gr.Checkbox(label="Randomize Seed", value=True)
+                    guidance_scale = gr.Slider(label="Guidance Scale", minimum=1.0, maximum=10.0, step=0.1, value=1.0)
+                    steps = gr.Slider(label="Inference Steps", minimum=1, maximum=50, step=1, value=4)
+                    target_megapixels = gr.Slider(
+                        label="Target Megapixels (canvas, 0 = match input area)",
+                        minimum=0.0,
+                        maximum=6.0,
+                        step=0.1,
+                        value=1.0,
+                    )
+                    resolution_multiple = gr.Dropdown(
+                        label="Resolution lattice multiple (anti-drift)",
+                        choices=[32, 56, 112],
+                        value=56,
+                        interactive=True,
+                    )
+                    vae_ref_megapixels = gr.Slider(
+                        label="Extra refs VAE megapixels override (0 = use canvas)",
+                        minimum=0.0,
+                        maximum=6.0,
+                        step=0.1,
+                        value=0.0,
+                    )
+                    decoder_vae = gr.Dropdown(
+                        label="Decoder VAE",
+                        choices=["qwen", "wan2x"],
+                        value="qwen",
+                        interactive=True,
+                    )
+                    keep_decoder_2x = gr.Checkbox(
+                        label="Keep 2× output (wan2x only)",
+                        value=False,
+                    )
+                    extras_condition_only = gr.Checkbox(
+                        label="Extra references are conditioning-only (exclude from VAE)",
+                        value=True,
+                    )
+                    pad_to_canvas = gr.Checkbox(
+                        label="Pad images to canvas aspect (avoid warping)",
+                        value=True,
+                    )
+
+                    # ✅ NEW: VAE tiling toggle (OFF by default)
+                    vae_tiling = gr.Checkbox(
+                        label="VAE tiling (lower VRAM, slower)",
+                        value=False,
+                    )
+
+        # On LoRA selection: preset prompt + toggle Image 2
+        lora_adapter.change(
+            fn=on_lora_change_ui,
+            inputs=[lora_adapter, prompt, extras_condition_only],
+            outputs=[prompt, input_image_2, extras_condition_only],
+        )
 
-    # LoRA change
-    lora_adapter.change(_on_lora_change, inputs=[lora_adapter], outputs=[prompt, img2, bfs_tip])
+        gr.Examples(
+            examples=[
+                ["examples/5.jpg", "Remove shadows and relight the image using soft lighting.", "Light-Restoration"],
+                ["examples/4.jpg", "Use a subtle golden-hour filter with smooth light diffusion.", "Relight"],
+                ["examples/2.jpeg", "Rotate the camera 45 degrees to the left.", "Multiple-Angles"],
+                [
+                    "examples/12.jpg",
+                    "flatcolor Desaturate the image and lower the contrast to create a flat, ungraded look similar to a camera log profile. Preserve details in the highlights and shadows.",
+                    "Flat-Log",
+                ],
+                ["examples/7.jpg", "Light source from the Right Rear", "Multi-Angle-Lighting"],
+                ["examples/10.jpeg", "Upscale the image.", "Upscale-Image"],
+                ["examples/7.jpg", "Light source from the Below", "Multi-Angle-Lighting"],
+                ["examples/2.jpeg", "Switch the camera to a top-down right corner view.", "Multiple-Angles"],
+                [
+                    "examples/9.jpg",
+                    "The camera moves slightly forward as sunlight breaks through the clouds, casting a soft glow around the character's silhouette in the mist. Realistic cinematic style, atmospheric depth.",
+                    "Next-Scene",
+                ],
+                ["examples/8.jpg", "Make the subjects skin details more prominent and natural.", "Edit-Skin"],
+                ["examples/6.jpg", "Switch the camera to a bottom-up view.", "Multiple-Angles"],
+                ["examples/6.jpg", "Rotate the camera 180 degrees upside down.", "Multiple-Angles"],
+                ["examples/4.jpg", "Rotate the camera 45 degrees to the right.", "Multiple-Angles"],
+                ["examples/4.jpg", "Switch the camera to a top-down view.", "Multiple-Angles"],
+                ["examples/4.jpg", "Switch the camera to a wide-angle lens.", "Multiple-Angles"],
+                ["examples/11.jpg", "Upscale this picture to 4K resolution.", "Upscale2K"],
+            ],
+            inputs=[input_image_1, prompt, lora_adapter],
+            outputs=[output_image, seed, last_output],
+            fn=infer_example,
+            cache_examples=False,
+            label="Examples",
+        )
 
-    # Run
-    run_btn.click(
-        infer,
+    run_button.click(
+        fn=infer,
         inputs=[
-            img1,
-            img2,
-            extra_gallery,
+            input_image_1,
+            input_image_2,
+            input_images_extra,
             prompt,
             lora_adapter,
             seed,
             randomize_seed,
-            guidance,
+            guidance_scale,
             steps,
             target_megapixels,
-            use_input_area,
-            keep_2x_output,
-            vae_tiling,
             extras_condition_only,
+            pad_to_canvas,
+            vae_tiling,
             resolution_multiple,
             vae_ref_megapixels,
-            use_depth,
-            derived_on_gpu,
+            decoder_vae,
+            keep_decoder_2x,
         ],
-        outputs=[out_img, seed, derived_preview],
+        outputs=[output_image, seed, last_output],
     )
 
     # 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])
+    btn_out_to_img1.click(fn=set_output_as_image1, inputs=[last_output], outputs=[input_image_1])
+    btn_out_to_img2.click(fn=set_output_as_image2, inputs=[last_output], outputs=[input_image_2])
+    btn_out_to_extra.click(fn=set_output_as_extra, inputs=[last_output, input_images_extra], outputs=[input_images_extra])
+
+    # Derived conditioning: append pose/depth map as extra ref (UI shows preview)
+    add_derived_btn.click(
+        fn=add_derived_ref,
+        inputs=[input_image_1, input_images_extra, derived_type, derived_use_gpu],
+        outputs=[input_images_extra, derived_preview],
+    )
 
-demo.queue(max_size=32).launch()
+if __name__ == "__main__":
+    demo.queue(max_size=30).launch(
+        css=css,
+        theme=orange_red_theme,
+        mcp_server=True,
+        ssr_mode=False,
+        show_error=True,
+    )