Update app.py

app.py

@@ -2,23 +2,23 @@ import os
 import re
 import gc
 import traceback
+import base64
+import io
 import gradio as gr
 import numpy as np
 import spaces
 import torch
 import random
-from PIL import Image, ImageDraw
-from typing import Iterable, Optional
+from PIL import Image
+from typing import Iterable, Optional, Tuple
 
 from transformers import (
-    AutoProcessor,
-    RTDetrForObjectDetection,
-    VitPoseForPoseEstimation,
     AutoImageProcessor,
     AutoModelForDepthEstimation,
 )
 
 from huggingface_hub import hf_hub_download
+from huggingface_hub import InferenceClient
 from safetensors.torch import load_file as safetensors_load_file
 
 from gradio.themes import Soft
@@ -137,7 +137,6 @@ 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
@@ -181,19 +180,15 @@ def _load_pipe_with_version(version: str) -> QwenImageEditPlusPipeline:
     return p
 
 
-# Forgiving load: try env/default version, fallback to v19 if it fails
 try:
     pipe = _load_pipe_with_version(AIO_VERSION)
-except Exception as e:
+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)
 
-# Apply FA3 Optimization
 try:
     pipe.transformer.set_attn_processor(QwenDoubleStreamAttnProcessorFA3())
     print("Flash Attention 3 Processor set successfully.")
@@ -203,56 +198,17 @@ except Exception as e:
 MAX_SEED = np.iinfo(np.int32).max
 
 # ============================================================
-# Derived conditioning (Transformers): Pose + Depth
+# Derived conditioning (Depth Anything) ONLY — ViTPose removed
 # ============================================================
-# Pose estimation uses ViTPose (top-down). Official docs show RT-DETR -> ViTPose flow:
-# https://huggingface.co/docs/transformers/model_doc/vitpose
-# Depth uses Depth Anything V2 Small (Transformers-compatible):
-# https://huggingface.co/depth-anything/Depth-Anything-V2-Small-hf
 
-POSE_MODEL_ID = "usyd-community/vitpose-base-simple"
-POSE_DETECTOR_ID = "PekingU/rtdetr_r50vd_coco_o365"
 DEPTH_MODEL_ID = "depth-anything/Depth-Anything-V2-Small-hf"
-
-# Lazy caches keyed by device string ("cpu" / "cuda")
-_POSE_CACHE = {}
 _DEPTH_CACHE = {}
 
-# COCO-17 skeleton connections (approx "OpenPose-like" stick figure)
-COCO17_EDGES = [
-    (0, 1), (0, 2), (1, 3), (2, 4),      # head
-    (5, 6),                              # shoulders
-    (5, 7), (7, 9),                      # left arm
-    (6, 8), (8, 10),                     # right arm
-    (5, 11), (6, 12), (11, 12),          # torso
-    (11, 13), (13, 15),                  # left leg
-    (12, 14), (14, 16),                  # right leg
-]
 
 def _derived_device(use_gpu: bool) -> torch.device:
     return torch.device("cuda" if (use_gpu and torch.cuda.is_available()) else "cpu")
 
 
-def _load_pose_models(dev: torch.device):
-    key = str(dev)
-    if key in _POSE_CACHE:
-        return _POSE_CACHE[key]
-
-    # Detector (optional but used for multi-person boxes)
-    det_proc = AutoProcessor.from_pretrained(POSE_DETECTOR_ID)
-    det_model = RTDetrForObjectDetection.from_pretrained(POSE_DETECTOR_ID).to(dev)
-
-    # Pose model
-    pose_proc = AutoProcessor.from_pretrained(POSE_MODEL_ID)
-    pose_model = VitPoseForPoseEstimation.from_pretrained(POSE_MODEL_ID).to(dev)
-
-    det_model.eval()
-    pose_model.eval()
-
-    _POSE_CACHE[key] = (det_proc, det_model, pose_proc, pose_model)
-    return _POSE_CACHE[key]
-
-
 def _load_depth_models(dev: torch.device):
     key = str(dev)
     if key in _DEPTH_CACHE:
@@ -266,116 +222,7 @@ def _load_depth_models(dev: torch.device):
     return _DEPTH_CACHE[key]
 
 
-def _draw_skeleton_on_blank(
-    size: tuple[int, int],
-    persons_keypoints: list[np.ndarray],
-    persons_scores: list[np.ndarray],
-    kp_thresh: float = 0.20,
-    point_r: int = 3,
-    line_w: int = 3,
-) -> Image.Image:
-    w, h = size
-    canvas = Image.new("RGB", (w, h), (0, 0, 0))
-    draw = ImageDraw.Draw(canvas)
-
-    for kps, sc in zip(persons_keypoints, persons_scores):
-        # Draw edges
-        for a, b in COCO17_EDGES:
-            if a >= len(sc) or b >= len(sc):
-                continue
-            if sc[a] < kp_thresh or sc[b] < kp_thresh:
-                continue
-            xa, ya = float(kps[a, 0]), float(kps[a, 1])
-            xb, yb = float(kps[b, 0]), float(kps[b, 1])
-            draw.line([(xa, ya), (xb, yb)], fill=(255, 255, 255), width=line_w)
-
-        # Draw keypoints
-        for i in range(min(len(sc), len(kps))):
-            if sc[i] < kp_thresh:
-                continue
-            x, y = float(kps[i, 0]), float(kps[i, 1])
-            draw.ellipse(
-                [(x - point_r, y - point_r), (x + point_r, y + point_r)],
-                fill=(255, 255, 255),
-                outline=None,
-            )
-
-    return canvas
-
-
-def make_pose_map(
-    img: Image.Image,
-    *,
-    use_gpu: bool,
-    mode: str,
-    det_thresh: float = 0.30,
-    max_people: int = 4,
-) -> Image.Image:
-    """Return an OpenPose-like skeleton map (RGB) using Transformers models.
-
-    mode:
-      - "fast": full-frame box (no detector). Good when Image 1 is already a single subject.
-      - "detect": RT-DETR person boxes -> ViTPose. Better for multi-person scenes.
-    """
-    img = img.convert("RGB")
-    dev = _derived_device(use_gpu)
-    det_proc, det_model, pose_proc, pose_model = _load_pose_models(dev)
-
-    w, h = img.size
-
-    if mode == "fast":
-        # Single box covering whole image, COCO format [x, y, w, h]
-        boxes = np.array([[0.0, 0.0, float(w), float(h)]], dtype=np.float32)
-    else:
-        # Detect people
-        inputs = det_proc(images=img, return_tensors="pt").to(dev)
-        with torch.no_grad():
-            outputs = det_model(**inputs)
-
-        results = det_proc.post_process_object_detection(
-            outputs,
-            target_sizes=torch.tensor([(h, w)], device=dev),
-            threshold=det_thresh,
-        )[0]
-
-        # COCO label 0 is "person" for COCO-trained detectors
-        person_boxes = results["boxes"][results["labels"] == 0].detach().cpu().numpy()
-
-        if person_boxes.size == 0:
-            # Fallback to full-frame
-            boxes = np.array([[0.0, 0.0, float(w), float(h)]], dtype=np.float32)
-        else:
-            # Convert VOC x1,y1,x2,y2 to COCO x,y,w,h
-            person_boxes[:, 2] = person_boxes[:, 2] - person_boxes[:, 0]
-            person_boxes[:, 3] = person_boxes[:, 3] - person_boxes[:, 1]
-            boxes = person_boxes.astype(np.float32)
-
-    if boxes.shape[0] > max_people:
-        boxes = boxes[:max_people]
-
-    pose_inputs = pose_proc(img, boxes=[boxes], return_tensors="pt").to(dev)
-    with torch.no_grad():
-        pose_outputs = pose_model(**pose_inputs)
-
-    pose_results = pose_proc.post_process_pose_estimation(pose_outputs, boxes=[boxes])[0]
-
-    persons_kps = []
-    persons_sc = []
-    for pr in pose_results:
-        kps = pr["keypoints"].detach().cpu().numpy()
-        sc = pr["scores"].detach().cpu().numpy()
-        persons_kps.append(kps)
-        persons_sc.append(sc)
-
-    if not persons_kps:
-        # No pose found; return black canvas
-        return Image.new("RGB", img.size, (0, 0, 0))
-
-    return _draw_skeleton_on_blank(img.size, persons_kps, persons_sc)
-
-
 def make_depth_map(img: Image.Image, *, use_gpu: bool) -> Image.Image:
-    """Return a grayscale (RGB) depth map using Depth Anything V2 Small."""
     img = img.convert("RGB")
     dev = _derived_device(use_gpu)
     proc, model = _load_depth_models(dev)
@@ -386,10 +233,7 @@ def make_depth_map(img: Image.Image, *, use_gpu: bool) -> Image.Image:
     with torch.no_grad():
         out = model(**inputs)
 
-    # predicted_depth: (B, H, W)
     pred = out.predicted_depth
-
-    # Upsample to original image size
     pred = torch.nn.functional.interpolate(
         pred.unsqueeze(1),
         size=(img.height, img.width),
@@ -403,8 +247,19 @@ def make_depth_map(img: Image.Image, *, use_gpu: bool) -> Image.Image:
     arr = arr / denom
 
     depth8 = (arr * 255.0).clip(0, 255).astype(np.uint8)
-    depth_img = Image.fromarray(depth8, mode="L").convert("RGB")
-    return depth_img
+    return Image.fromarray(depth8, mode="L").convert("RGB")
+
+
+def _to_pil_rgb(item):
+    if item is None:
+        return None
+    if isinstance(item, (tuple, list)) and len(item) >= 1:
+        item = item[0]
+    if isinstance(item, Image.Image):
+        return item.convert("RGB")
+    if isinstance(item, np.ndarray):
+        return Image.fromarray(item).convert("RGB")
+    return None
 
 
 def _append_to_gallery(existing, new_img: Image.Image):
@@ -417,6 +272,7 @@ def _append_to_gallery(existing, new_img: Image.Image):
     items.append(new_img)
     return items
 
+
 # ============================================================
 # LoRA adapters + presets
 # ============================================================
@@ -486,7 +342,7 @@ ADAPTER_SPECS = {
         "weights": "bfs_head_v5_2511_original.safetensors",
         "adapter_name": "BFS-Best-Faceswap",
         "strength": 1.0,
-        "needs_alpha_fix": True,  # <-- fixes KeyError 'img_in.alpha'
+        "needs_alpha_fix": True,
     },
     "BFS-Best-FaceSwap-merge": {
         "type": "single",
@@ -496,7 +352,7 @@ ADAPTER_SPECS = {
         "weights": "bfs_head_v5_2511_merged_version_rank_32_fp32.safetensors",
         "adapter_name": "BFS-Best-Faceswap-merge",
         "strength": 1.1,
-        "needs_alpha_fix": True,  # <-- fixes KeyError 'img_in.alpha'
+        "needs_alpha_fix": True,
     },
     "F2P": {
         "type": "single",
@@ -575,480 +431,507 @@ 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. 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.",
+    "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 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, 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:
     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]:
-    """Compute (width, height) that matches image aspect ratio and approximates target_area.
+    multiple_of: int = 64,
+) -> Tuple[int, int]:
+    w0, h0 = image.size
+    if w0 <= 0 or h0 <= 0:
+        return 512, 512
+    aspect = w0 / h0
+    w = int((target_area * aspect) ** 0.5)
+    h = int(w / aspect) if aspect != 0 else int((target_area) ** 0.5)
+    w = _round_to_multiple(w, multiple_of)
+    h = _round_to_multiple(h, multiple_of)
+    w = max(multiple_of, w)
+    h = max(multiple_of, h)
+    return w, h
+
+
+def get_target_area_for_lora(image: Image.Image, lora_adapter: str, target_megapixels: float) -> int:
+    spec = ADAPTER_SPECS.get(lora_adapter, {})
+    long_edge = spec.get("target_long_edge", None)
 
-    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
+    if long_edge:
+        w0, h0 = image.size
+        if w0 <= 0 or h0 <= 0:
+            return int(1.0 * 1024 * 1024)
+        scale = float(long_edge) / float(max(w0, h0))
+        w = int(w0 * scale)
+        h = int(h0 * scale)
+        return max(64 * 64, w * h)
 
-    # Use the pipeline's own area->(w,h) helper for consistency.
-    from qwenimage.pipeline_qwenimage_edit_plus import calculate_dimensions
+    mp = float(target_megapixels)
+    return max(64 * 64, int(mp * 1_000_000))
 
-    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:
-    """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, {})
+# ============================================================
+# Helpers: LoRA loading + alpha fix
+# ============================================================
 
-    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
+def _download_from_hf(repo_id: str, filename: str) -> str:
+    return hf_hub_download(repo_id=repo_id, filename=filename)
 
-    # Legacy support (e.g. Upscale2K)
-    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
 
-    # User default
-    return int(float(user_target_megapixels) * 1024 * 1024)
+def _maybe_apply_alpha_fix(state_dict: dict) -> dict:
+    if "img_in.alpha" not in state_dict:
+        for k in list(state_dict.keys()):
+            if k.endswith("img_in.weight") or k.endswith("img_in.bias"):
+                t = state_dict[k]
+                if hasattr(t, "new_zeros"):
+                    state_dict["img_in.alpha"] = t.new_zeros(())
+                break
+    return state_dict
 
-# ============================================================
-# Helpers: multi-input routing + gallery normalization
-# ============================================================
 
+def _load_single_lora(spec: dict):
+    local_path = _download_from_hf(spec["repo"], spec["weights"])
+    sd = safetensors_load_file(local_path)
+    if spec.get("needs_alpha_fix", False):
+        sd = _maybe_apply_alpha_fix(sd)
+    pipe.load_lora_weights(sd, adapter_name=spec["adapter_name"])
+    LOADED_ADAPTERS.add(spec["adapter_name"])
 
-def lora_requires_two_images(lora_adapter: str) -> bool:
-    return bool(ADAPTER_SPECS.get(lora_adapter, {}).get("requires_two_images", False))
 
+def _ensure_loaded_and_get_active_adapters(lora_adapter: str):
+    spec = ADAPTER_SPECS.get(lora_adapter, None)
+    if spec is None:
+        return [], []
+
+    if spec["type"] == "single":
+        if spec["adapter_name"] not in LOADED_ADAPTERS:
+            _load_single_lora(spec)
+        return [spec["adapter_name"]], [spec.get("strength", 1.0)]
 
-def image2_label_for_lora(lora_adapter: str) -> str:
-    return str(ADAPTER_SPECS.get(lora_adapter, {}).get("image2_label", "Upload Reference (Image 2)"))
+    adapter_names = []
+    weights = []
+    for part in spec["parts"]:
+        if part["adapter_name"] not in LOADED_ADAPTERS:
+            _load_single_lora(part)
+        adapter_names.append(part["adapter_name"])
+        weights.append(part.get("strength", 1.0))
+    return adapter_names, weights
 
 
-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
+def lora_requires_two_images(lora_adapter: str) -> bool:
+    spec = ADAPTER_SPECS.get(lora_adapter, {})
+    return bool(spec.get("requires_two_images", False))
 
-    # Gallery often returns (image, caption)
-    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")
+def get_image2_label_for_lora(lora_adapter: str) -> str:
+    spec = ADAPTER_SPECS.get(lora_adapter, {})
+    return spec.get("image2_label", "Upload Reference (Image 2)")
 
-    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 build_labeled_images(img1: Image.Image, img2: Optional[Image.Image], extras: list[Image.Image]):
+    labeled = {"image_1": img1}
+    if img2 is not None:
+        labeled["image_2"] = img2
+    for ex in extras:
+        labeled[f"image_{len(labeled) + 1}"] = ex
+    return labeled
 
 
-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
+# ============================================================
+# UI: lora change handler
+# ============================================================
 
-    labeled[f"image_{idx}"] = img1
-    idx += 1
 
-    if img2 is not None:
-        labeled[f"image_{idx}"] = img2
-        idx += 1
+def on_lora_change_ui(lora_adapter, current_prompt, current_extras_condition_only):
+    preset = LORA_PRESET_PROMPTS.get(lora_adapter, None)
+    prompt_update = gr.update(value=preset) if preset else gr.update(value=current_prompt)
 
-    if extra_imgs:
-        for im in extra_imgs:
-            if im is None:
-                continue
-            labeled[f"image_{idx}"] = im
-            idx += 1
+    needs_two = lora_requires_two_images(lora_adapter)
+    img2_update = gr.update(visible=needs_two, label=get_image2_label_for_lora(lora_adapter))
 
-    return labeled
+    extras_update = gr.update(value=True) if needs_two else gr.update(value=current_extras_condition_only)
+    return prompt_update, img2_update, extras_update
 
 
 # ============================================================
-# Helpers: BFS alpha key fix
+# Output routing + derived conditioning
 # ============================================================
 
 
-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).
+def set_output_as_image1(last):
+    if last is None:
+        raise gr.Error("No output available yet.")
+    return gr.update(value=last)
 
-    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
-        if k.endswith(".lora_down.weight") and v.ndim >= 1:
-            base = k[: -len(".lora_down.weight")]
-            rank = int(v.shape[0])
-            bases[base] = rank
+def set_output_as_image2(last):
+    if last is None:
+        raise gr.Error("No output available yet.")
+    return gr.update(value=last)
 
-    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
+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)
 
-        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
+@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.")
 
+    if derived_type == "None":
+        return gr.update(value=existing_extra), gr.update(visible=False, value=None)
 
-def _filter_to_diffusers_lora_keys(state_dict: dict) -> tuple[dict, dict]:
-    """Return (filtered_state_dict, stats).
+    base = img1.convert("RGB")
 
-    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...`).
+    if derived_type == "Depth (Depth Anything V2 Small)":
+        derived = make_depth_map(base, use_gpu=bool(derived_use_gpu))
+    else:
+        raise gr.Error(f"Unknown derived type: {derived_type}")
 
-    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`)
+    new_gallery = _append_to_gallery(existing_extra, derived)
+    return gr.update(value=new_gallery), gr.update(visible=True, value=derived)
 
-    It also drops known patch keys (`*.diff`, `*.diff_b`) and everything else.
-    """
 
-    keep_suffixes = (
-        ".lora_up.weight",
-        ".lora_down.weight",
-        ".lora_mid.weight",
-        ".alpha",
-        ".lora_alpha",
-    )
+# ============================================================
+# Prompt Helper (outsourced VLM calls, UI stays clean)
+# ============================================================
 
-    dropped_patch = 0
-    dropped_other = 0
-    kept = 0
-    normalized_alpha = 0
-
-    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
-
-    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:
-    """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.
-    """
+# Configuration via env vars (no UI clutter)
+HF_TOKEN = os.environ.get("HF_TOKEN", "").strip() or os.environ.get("HUGGINGFACEHUB_API_TOKEN", "").strip()
+HF_PROVIDER = os.environ.get("HF_PROVIDER", "nebius").strip()
+HF_VLM_MODEL = os.environ.get("HF_VLM_MODEL", "Qwen/Qwen2.5-VL-7B-Instruct").strip()
 
-    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
+_client_cache = {}
 
 
-def _load_lora_weights_with_fallback(repo: str, weight_name: str, adapter_name: str, needs_alpha_fix: bool = False):
+def _get_client() -> InferenceClient:
+    key = (HF_PROVIDER, bool(HF_TOKEN))
+    if key in _client_cache:
+        return _client_cache[key]
+    if not HF_TOKEN:
+        raise gr.Error("Captioning is not configured (missing HF_TOKEN).")
+    client = InferenceClient(provider=HF_PROVIDER, api_key=HF_TOKEN)
+    _client_cache[key] = client
+    return client
+
+
+def _encode_image_data_url(img: Image.Image, max_side: int = 1536, fmt: str = "PNG") -> str:
     """
-    Normal path: pipe.load_lora_weights(repo, weight_name=..., adapter_name=...)
-    BFS fallback: download safetensors, inject missing alpha keys, then load from dict.
+    Converts PIL to data URL (base64). Downscales to keep payload reasonable.
     """
-    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
-
-        print(
-            "⚠️ LoRA load failed (will try safe dict fallback). "
-            f"Adapter={adapter_name!r} file={weight_name!r} error={type(e).__name__}: {e}"
-        )
+    img = img.convert("RGB")
+    w, h = img.size
+    scale = min(1.0, float(max_side) / float(max(w, h))) if max(w, h) > 0 else 1.0
+    if scale < 1.0:
+        img = img.resize((max(1, int(w * scale)), max(1, int(h * scale))), Image.LANCZOS)
 
-        local_path = hf_hub_download(repo_id=repo, filename=weight_name)
-        sd = safetensors_load_file(local_path)
+    buf = io.BytesIO()
+    img.save(buf, format=fmt)
+    b64 = base64.b64encode(buf.getvalue()).decode("utf-8")
+    mime = "image/png" if fmt.upper() == "PNG" else "image/jpeg"
+    return f"data:{mime};base64,{b64}"
 
-        # 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)
+def _chat_with_image(
+    system_prompt: str,
+    user_text: str,
+    image: Image.Image,
+    *,
+    max_tokens: int,
+    temperature: float,
+) -> str:
+    client = _get_client()
+    data_url = _encode_image_data_url(image)
+
+    messages = [
+        {"role": "system", "content": system_prompt},
+        {
+            "role": "user",
+            "content": [
+                {"type": "text", "text": user_text},
+                {"type": "image_url", "image_url": {"url": data_url}},
+            ],
+        },
+    ]
+
+    # Hugging Face chat.completions interface
+    resp = client.chat.completions.create(
+        model=HF_VLM_MODEL,
+        messages=messages,
+        max_tokens=int(max_tokens),
+        temperature=float(temperature),
+    )
+    return (resp.choices[0].message.content or "").strip()
 
-        # 3) Duplicate stripped keys (remove `diffusion_model.`) for compatibility.
-        sd = _duplicate_stripped_prefix_keys(sd)
 
-        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']}"
-        )
+def _chat_text_only(
+    system_prompt: str,
+    user_text: str,
+    *,
+    max_tokens: int,
+    temperature: float,
+) -> str:
+    client = _get_client()
+    messages = [
+        {"role": "system", "content": system_prompt},
+        {"role": "user", "content": [{"type": "text", "text": user_text}]},
+    ]
+    resp = client.chat.completions.create(
+        model=HF_VLM_MODEL,
+        messages=messages,
+        max_tokens=int(max_tokens),
+        temperature=float(temperature),
+    )
+    return (resp.choices[0].message.content or "").strip()
 
-        pipe.load_lora_weights(sd, adapter_name=adapter_name)
-        return
 
+def _has_header(text: str, header: str) -> bool:
+    return header in (text or "")
 
-# ============================================================
-# LoRA loader: single/package + strengths
-# ============================================================
 
+def _enforce_once_retry_image(system_prompt: str, user_text: str, image: Image.Image, header: str, max_tokens: int, temperature: float) -> str:
+    out = _chat_with_image(system_prompt, user_text, image, max_tokens=max_tokens, temperature=temperature)
+    if _has_header(out, header):
+        return out
 
-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}")
+    # one strict retry
+    retry_user = (
+        user_text
+        + "\n\nIMPORTANT: You did not follow the required output format. "
+        + f"Return EXACTLY the block starting with {header} and fill each line. No extra text."
+    )
+    out2 = _chat_with_image(system_prompt, retry_user, image, max_tokens=max_tokens, temperature=temperature)
+    return out2
 
-    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"]
-            adapter_name = part["adapter_name"]
-            strength = float(part.get("strength", 1.0))
-            needs_alpha_fix = bool(part.get("needs_alpha_fix", False))
-
-            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(adapter_name)
-            adapter_weights.append(strength)
+def _enforce_once_retry_text(system_prompt: str, user_text: str, header: str, max_tokens: int, temperature: float) -> str:
+    out = _chat_text_only(system_prompt, user_text, max_tokens=max_tokens, temperature=temperature)
+    if _has_header(out, header):
+        return out
 
-    else:
-        repo = spec["repo"]
-        weights = spec["weights"]
-        adapter_name = spec["adapter_name"]
-        strength = float(spec.get("strength", 1.0))
-        needs_alpha_fix = bool(spec.get("needs_alpha_fix", False))
-
-        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. ---")
+    retry_user = (
+        user_text
+        + "\n\nIMPORTANT: You did not follow the required output format. "
+        + f"Return EXACTLY the sections starting with {header}. No extra text."
+    )
+    return _chat_text_only(system_prompt, retry_user, max_tokens=max_tokens, temperature=temperature)
+
+
+# --------- BASE (Pic1) extraction prompt (no identity) ----------
+BFS_BASE_SYSTEM = """You are extracting non-identity facial and contextual signals from Picture 1 (BASE) for a head/face swap.
+
+CRITICAL: DO NOT describe identity/likeness traits. That means:
+- No age, ethnicity/race/nationality guesses, attractiveness judgments, “looks like X”
+- No skin tone, facial structure descriptions, “round face”, “strong jaw”, etc.
+- No hair color/style as identity markers (only mention hair if it occludes the face, e.g. “hair covering left eye”)
+
+Focus ONLY on:
+- Head pose (yaw/pitch/roll, tilt, chin/jaw position)
+- Gaze and eyelids (direction, openness)
+- Micro-expressions / muscle cues (brow knit/raise, squint, lip tension, mouth corners, cheek tension, jaw set)
+- Mouth details (open/closed, teeth, tongue if visible)
+- Mood inference (max 2 labels) with visible evidence cues
+- Occlusions and interactions (hands, objects, glasses, shadows) relevant to face recreation
+- Visibility notes (unclear/occluded/shadowed)
+
+Output format (return exactly this block, nothing else):
+
+[BASE_SIGNALS_PIC1]
+Head pose:
+Gaze & eyelids:
+Expression (muscle cues):
+Mouth details:
+Mood (max 2 labels):
+Evidence for mood (visible cues only):
+Occlusions & interactions:
+Visibility notes (unclear/occluded/shadowed areas):
+"""
 
-        adapter_names = [adapter_name]
-        adapter_weights = [strength]
+BFS_BASE_USER = """Analyze the single provided image as Picture 1 (BASE).
+Fill every line with either an observation or the word "unclear". Keep it concise."""
+
+# --------- DONOR (Pic2) extraction prompt (identity only) ----------
+BFS_DONOR_SYSTEM = """You are extracting inherent identity/likeness traits from Picture 2 (DONOR) for a head/face swap.
+
+CRITICAL: DO NOT describe expression, mood, gaze direction, head pose/rotation, body pose, or actions.
+
+Focus ONLY on visible physical traits:
+- Face shape & proportions (jawline, cheekbones, chin shape)
+- Skin tone/undertone + texture (freckles/moles only if visible)
+- Eyes (color, shape), brows (shape/thickness)
+- Nose structure (bridge, tip, nostrils)
+- Lips/mouth shape (fullness, cupid’s bow)
+- Chin/jaw details
+- Hair (color, style, hairline)
+- Distinctive traits (scars/moles/freckles if visible)
+- Visibility notes (unclear/occluded/shadowed)
+
+Output format (return exactly this block, nothing else):
+
+[DONOR_TRAITS_PIC2]
+Face shape & proportions:
+Skin tone & texture:
+Eyes & brows:
+Nose structure:
+Lips & mouth shape:
+Chin/jaw details:
+Hair (color, style, hairline):
+Distinctive traits (scars/moles/freckles if visible):
+Visibility notes (unclear/occluded/shadowed areas):
+"""
 
-    return adapter_names, adapter_weights
+BFS_DONOR_USER = """Analyze the single provided image as Picture 2 (DONOR).
+Fill every line with either an observation or the word "unclear". Keep it concise."""
 
+# --------- Text-only prompt builder ----------
+BFS_BUILDER_SYSTEM = """You are a prompt editor for BFS-BestFaceSwap.
 
-# ============================================================
-# UI handlers
-# ============================================================
+Input you may receive:
+- A core prompt (already includes head_swap instructions)
+- BASE_SIGNALS_PIC1 text (pose/expression/mood/occlusions; non-identity)
+- Optional DONOR_TRAITS_PIC2 text (identity-only traits)
 
+Your job:
+- Produce a compact addendum that improves expressiveness transfer and reduces ambiguity.
+- Do NOT add any identity traits from the base signals.
+- Do NOT add any pose/expression/mood from donor traits.
+- Prefer concrete, visible cues over vague adjectives.
+- Keep it short (ideally 6–14 lines total).
+- If donor traits are missing or mostly "unclear", omit donor section entirely.
 
-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)
+Output EXACTLY two sections (donor section may be omitted if not provided/usable):
+[ADDENDUM_BASE]
+(bullets or short lines; use the best cues from BASE_SIGNALS)
 
-    # 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)")
+[ADDENDUM_DONOR]
+(optional; only if donor traits contain useful visible info; no pose/expression)
+"""
 
-    # 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
-# ============================================================
-# UI helpers: output routing + derived conditioning
-# ============================================================
+def scrub_placeholder(text: str, enabled: bool) -> str:
+    # Placeholder for future strict scrubber pass (no-op).
+    return text
 
-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 caption_base_pic1(
+    img1,
+    max_new_tokens: int,
+    temperature: float,
+    strict_scrubber: bool,
+    show_debug: bool,
+):
+    if img1 is None:
+        raise gr.Error("Please upload Image 1 (base) first.")
 
-def set_output_as_image2(last):
-    if last is None:
-        raise gr.Error("No output available yet.")
-    return gr.update(value=last)
+    raw = _enforce_once_retry_image(
+        BFS_BASE_SYSTEM,
+        BFS_BASE_USER,
+        img1,
+        header="[BASE_SIGNALS_PIC1]",
+        max_tokens=int(max_new_tokens),
+        temperature=float(temperature),
+    )
+    out = scrub_placeholder(raw, enabled=bool(strict_scrubber))
+    debug = raw if bool(show_debug) else ""
+    return out, debug
 
 
-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)
+@spaces.GPU
+def caption_donor_pic2(
+    img2,
+    max_new_tokens: int,
+    temperature: float,
+    strict_scrubber: bool,
+    show_debug: bool,
+):
+    if img2 is None:
+        raise gr.Error("Please upload Image 2 (donor) first.")
+
+    raw = _enforce_once_retry_image(
+        BFS_DONOR_SYSTEM,
+        BFS_DONOR_USER,
+        img2,
+        header="[DONOR_TRAITS_PIC2]",
+        max_tokens=int(max_new_tokens),
+        temperature=float(temperature),
+    )
+    out = scrub_placeholder(raw, enabled=bool(strict_scrubber))
+    debug = raw if bool(show_debug) else ""
+    return out, debug
 
 
-@spaces.GPU
-def add_derived_ref(img1, existing_extra, derived_type, derived_use_gpu, derived_max_people):
-    if img1 is None:
-        raise gr.Error("Please upload Image 1 first.")
+def _compose_final_prompt(core_prompt: str, addendum_text: str, mode: str) -> str:
+    core = (core_prompt or "").strip()
+    addendum = (addendum_text or "").strip()
+    if not addendum:
+        return core
 
-    if derived_type == "None":
-        return gr.update(value=existing_extra), gr.update(visible=False, value=None)
+    if (mode or "").lower().startswith("inject"):
+        injected = core
+        if "{BFS_ADDENDUM}" in injected:
+            injected = injected.replace("{BFS_ADDENDUM}", addendum + "\n")
+            return injected.strip()
 
-    base = img1.convert("RGB")
+    return (core + "\n\n" + addendum).strip()
 
-    if derived_type == "Pose (ViTPose, fast)":
-        derived = make_pose_map(base, use_gpu=bool(derived_use_gpu), mode="fast")
-    elif derived_type == "Pose (ViTPose + RT-DETR detect)":
-        derived = make_pose_map(
-            base,
-            use_gpu=bool(derived_use_gpu),
-            mode="detect",
-            max_people=int(derived_max_people),
-        )
-    elif derived_type == "Depth (Depth Anything V2 Small)":
-        derived = make_depth_map(base, use_gpu=bool(derived_use_gpu))
-    else:
-        raise gr.Error(f"Unknown derived type: {derived_type}")
 
-    new_gallery = _append_to_gallery(existing_extra, derived)
-    return gr.update(value=new_gallery), gr.update(visible=True, value=derived)
+@spaces.GPU
+def build_bfs_addendum_and_final_prompt(
+    core_prompt: str,
+    base_caption: str,
+    donor_caption: str,
+    integration_mode: str,
+    max_new_tokens: int,
+    temperature: float,
+    show_debug: bool,
+):
+    base = (base_caption or "").strip()
+    donor = (donor_caption or "").strip()
+    core = (core_prompt or "").strip()
+
+    if not base:
+        raise gr.Error("Generate BASE signals (Pic1) first (or paste them) before building an addendum.")
+
+    user_text = (
+        "CORE PROMPT:\n"
+        f"{core}\n\n"
+        "BASE_SIGNALS_PIC1:\n"
+        f"{base}\n\n"
+        "DONOR_TRAITS_PIC2:\n"
+        f"{donor if donor else '(none)'}\n\n"
+        "Produce the addendum now."
+    )
 
+    raw = _enforce_once_retry_text(
+        BFS_BUILDER_SYSTEM,
+        user_text,
+        header="[ADDENDUM_BASE]",
+        max_tokens=int(max_new_tokens),
+        temperature=float(temperature),
+    )
+
+    final_prompt = _compose_final_prompt(core, raw, integration_mode)
+    debug = raw if bool(show_debug) else ""
+    return raw, final_prompt, debug
 
 
 # ============================================================
@@ -1060,7 +943,7 @@ def add_derived_ref(img1, existing_extra, derived_type, derived_use_gpu, derived
 def infer(
     input_image_1,
     input_image_2,
-    input_images_extra,  # gallery multi-image box
+    input_images_extra,
     prompt,
     lora_adapter,
     seed,
@@ -1079,7 +962,6 @@ def infer(
     if input_image_1 is None:
         raise gr.Error("Please upload Image 1.")
 
-    # Handle "None"
     if lora_adapter == NONE_LORA:
         try:
             pipe.set_adapters([], adapter_weights=[])
@@ -1102,7 +984,6 @@ def infer(
     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:
@@ -1110,20 +991,14 @@ def infer(
             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,
@@ -1131,19 +1006,12 @@ def infer(
         multiple_of=int(pipe.vae_scale_factor * 2),
     )
 
-    # 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]
 
     try:
-        print(
-             "[DEBUG][infer] submitting request | "
-             f"lora_adapter={lora_adapter!r} seed={seed} prompt={prompt!r}"
-        )
-        
         result = pipe(
             image=pipe_images,
             prompt=prompt,
@@ -1170,8 +1038,20 @@ def infer_example(input_image, prompt, lora_adapter):
     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.
-    result, seed, last = infer(input_pil, None, None, prompt, lora_adapter, 0, True, guidance_scale, steps, 1.0, True, True)
+    result, seed, last = infer(
+        input_pil,
+        None,
+        None,
+        prompt,
+        lora_adapter,
+        0,
+        True,
+        guidance_scale,
+        steps,
+        1.0,
+        True,
+        True,
+    )
     return result, seed, last
 
 
@@ -1180,11 +1060,8 @@ def infer_example(input_image, prompt, lora_adapter):
 # ============================================================
 
 css = """
-#col-container {
-    margin: 0 auto;
-    max-width: 960px;
-}
-#main-title h1 {font-size: 2.1em !important;}
+#col-container { margin: 0 auto; max-width: 960px; }
+#main-title h1 { font-size: 2.1em !important; }
 """
 
 aio_status_line = (
@@ -1198,7 +1075,7 @@ with gr.Blocks() as demo:
         gr.Markdown(
             "Perform diverse image edits using specialized "
             "[LoRA](https://huggingface.co/models?other=base_model:adapter:Qwen/Qwen-Image-Edit-2511) adapters for the "
-            "[Qwen-Image-Edit](https://huggingface.co/Qwen/Qwen-Image-Edit-2511) model. Uses a Diffusers compatible extraction of the transformers from Phr00t's Rapid AIO merge. If a different AIO version is desired, copy the space and set the space variable to change version.'"
+            "[Qwen-Image-Edit](https://huggingface.co/Qwen/Qwen-Image-Edit-2511) model."
         )
         gr.Markdown(aio_status_line)
 
@@ -1222,11 +1099,45 @@ with gr.Blocks() as demo:
                     placeholder="e.g., transform into photo..",
                 )
 
+                with gr.Accordion("BFS Prompt Helper", open=False):
+                    with gr.Row():
+                        helper_max_tokens = gr.Slider(label="Max new tokens", minimum=64, maximum=1024, step=16, value=384)
+                        helper_temperature = gr.Slider(label="Temperature (0 = deterministic)", minimum=0.0, maximum=1.2, step=0.05, value=0.2)
+
+                    with gr.Row():
+                        strict_scrubber = gr.Checkbox(label="Strict scrubber (placeholder, no-op)", value=False)
+                        show_debug = gr.Checkbox(label="Show debug outputs", value=False)
+
+                    with gr.Row():
+                        btn_cap_base = gr.Button("Generate BASE signals (Pic1)", variant="secondary")
+                        btn_cap_donor = gr.Button("Generate DONOR traits (Pic2) (optional)", variant="secondary")
+
+                    with gr.Row():
+                        caption_pic1 = gr.Textbox(label="BASE signals (from Image 1)", lines=12, value="")
+                        caption_pic2 = gr.Textbox(label="DONOR traits (from Image 2) (optional)", lines=12, value="")
+
+                    with gr.Row():
+                        debug_base = gr.Textbox(label="Debug: raw BASE output", lines=8, visible=False)
+                        debug_donor = gr.Textbox(label="Debug: raw DONOR output", lines=8, visible=False)
+
+                    integration_mode = gr.Radio(
+                        label="How to apply addendum to the core prompt",
+                        choices=["Concatenate", "Inject (placeholder {BFS_ADDENDUM})"],
+                        value="Concatenate",
+                    )
+
+                    with gr.Row():
+                        btn_build_addendum = gr.Button("Build addendum + final prompt", variant="primary")
+                        btn_apply_final = gr.Button("Apply final prompt → Edit Prompt", variant="secondary")
+
+                    bfs_addendum = gr.Textbox(label="Built addendum (editable)", lines=10, value="")
+                    bfs_final_prompt = gr.Textbox(label="Final prompt preview (editable)", lines=10, value="")
+                    debug_builder = gr.Textbox(label="Debug: raw builder output", lines=8, visible=False)
+
                 run_button = gr.Button("Edit Image", variant="primary")
 
             with gr.Column():
                 output_image = gr.Image(label="Output Image", interactive=False, format="png", height=353)
-
                 last_output = gr.State(value=None)
 
                 with gr.Row():
@@ -1251,25 +1162,13 @@ with gr.Blocks() as demo:
                     )
 
                 with gr.Accordion("Advanced Settings", open=False, visible=True):
-                    with gr.Accordion("Derived Conditioning (Pose / Depth)", open=False):
+                    with gr.Accordion("Derived Conditioning (Depth)", open=False):
                         derived_type = gr.Dropdown(
                             label="Derived Type (from Image 1)",
-                            choices=[
-                                "None",
-                                "Pose (ViTPose, fast)",
-                                "Pose (ViTPose + RT-DETR detect)",
-                                "Depth (Depth Anything V2 Small)",
-                            ],
+                            choices=["None", "Depth (Depth Anything V2 Small)"],
                             value="None",
                         )
                         derived_use_gpu = gr.Checkbox(label="Use GPU for derived model", value=False)
-                        derived_max_people = gr.Slider(
-                            label="Max people (pose detect mode)",
-                            minimum=1,
-                            maximum=10,
-                            step=1,
-                            value=4,
-                        )
                         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)
@@ -1292,38 +1191,64 @@ with gr.Blocks() as demo:
                         value=True,
                     )
 
-        # On LoRA selection: preset prompt + toggle Image 2
+        # 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],
         )
 
+        # Debug visibility toggles
+        show_debug.change(
+            fn=lambda x: (
+                gr.update(visible=bool(x)),
+                gr.update(visible=bool(x)),
+                gr.update(visible=bool(x)),
+            ),
+            inputs=[show_debug],
+            outputs=[debug_base, debug_donor, debug_builder],
+        )
+
+        # Caption buttons (single-image)
+        btn_cap_base.click(
+            fn=caption_base_pic1,
+            inputs=[input_image_1, helper_max_tokens, helper_temperature, strict_scrubber, show_debug],
+            outputs=[caption_pic1, debug_base],
+        )
+
+        btn_cap_donor.click(
+            fn=caption_donor_pic2,
+            inputs=[input_image_2, helper_max_tokens, helper_temperature, strict_scrubber, show_debug],
+            outputs=[caption_pic2, debug_donor],
+        )
+
+        # Builder (text-only)
+        btn_build_addendum.click(
+            fn=build_bfs_addendum_and_final_prompt,
+            inputs=[
+                prompt,
+                caption_pic1,
+                caption_pic2,
+                integration_mode,
+                helper_max_tokens,
+                helper_temperature,
+                show_debug,
+            ],
+            outputs=[bfs_addendum, bfs_final_prompt, debug_builder],
+        )
+
+        # Apply final prompt to the Edit Prompt box
+        btn_apply_final.click(
+            fn=lambda x: gr.update(value=x),
+            inputs=[bfs_final_prompt],
+            outputs=[prompt],
+        )
+
         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],
@@ -1352,18 +1277,18 @@ with gr.Blocks() as demo:
         outputs=[output_image, seed, last_output],
     )
 
-    # Output routing buttons
+    # Output routing
     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)
+
+    # Derived conditioning: append depth map
     add_derived_btn.click(
         fn=add_derived_ref,
-        inputs=[input_image_1, input_images_extra, derived_type, derived_use_gpu, derived_max_people],
+        inputs=[input_image_1, input_images_extra, derived_type, derived_use_gpu],
         outputs=[input_images_extra, derived_preview],
     )
-    
+
 if __name__ == "__main__":
     demo.queue(max_size=30).launch(
         css=css,