Added new functionalities

Added controlnet-style depthmap and pose generator. Added bulk processing and multi-generations. Added canvas control. Added multiple QoL improvements

app.py

@@ -1,20 +1,57 @@
 import os
 import gc
+import csv
+import time
+import random
+import uuid
+import zipfile
+import threading
+from typing import Iterable
+
 import gradio as gr
 import numpy as np
 import spaces
 import torch
-import random
-import uuid
-from PIL import Image, ImageOps
-from typing import Iterable
-import time
-import threading
+from PIL import Image
+
 from logging_utils import log_inference
 
-# ── requirements.txt should include: ─────────────────────────────────────────
-#   spandrel
-# ─────────────────────────────────────────────────────────────────────────────
+from image_utils import (
+    fix_orientation,
+    compute_base_dimensions,
+    compute_canvas_dimensions,
+    fit_to_canvas,
+    on_base_image_change,
+    on_reference_change,
+    process_images,
+    reencode_upload,
+    send_editor_to_base,
+    send_editor_to_reference,
+    load_heic_to_editor,
+    send_output_to_base,
+    send_output_to_reference,
+    save_with_metadata,
+    build_pnginfo,
+    push_pil_to_base,
+    push_pil_to_reference,
+)
+
+from control_tools import (
+    generate_depthmap,
+    detect_pose,
+    render_pose_skeleton,
+    render_pose_overlay,
+    move_joint,
+    hide_joint,
+    clear_all_joints,
+    default_pose_template,
+    person_choices,
+    parse_person_idx,
+    joint_name_to_index,
+    OPENPOSE_KEYPOINT_NAMES,
+)
+
+# requirements.txt: spandrel, pillow-heif
 
 MODEL_VARIANT = os.environ.get("MODEL_VARIANT", "9B")
 if MODEL_VARIANT == "9B-KV":
@@ -23,11 +60,11 @@ if MODEL_VARIANT == "9B-KV":
 else:
     from diffusers import Flux2KleinPipeline as _PipeClass
     _MODEL_REPO = "black-forest-labs/FLUX.2-klein-9B"
-    MODEL_VARIANT = "9B"  # normalise any typos back to default
-from huggingface_hub import hf_hub_download
+    MODEL_VARIANT = "9B"
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
+# ── Theme ────────────────────────────────────────────────────────────────────
 from gradio.themes import Soft
 from gradio.themes.utils import colors, fonts, sizes
 
@@ -38,21 +75,13 @@ colors.orange_red = colors.Color(
 )
 
 class OrangeRedTheme(Soft):
-    def __init__(
-        self, *, primary_hue: colors.Color | str = colors.gray,
-        secondary_hue: colors.Color | str = colors.orange_red,
-        neutral_hue: colors.Color | str = colors.slate, text_size: sizes.Size | str = sizes.text_lg,
-        font: fonts.Font | str | Iterable[fonts.Font | str] = (
-            fonts.GoogleFont("Outfit"), "Arial", "sans-serif",
-        ),
-        font_mono: fonts.Font | str | Iterable[fonts.Font | str] = (
-            fonts.GoogleFont("IBM Plex Mono"), "ui-monospace", "monospace",
-        ),
-    ):
-        super().__init__(
-            primary_hue=primary_hue, secondary_hue=secondary_hue, neutral_hue=neutral_hue,
-            text_size=text_size, font=font, font_mono=font_mono,
-        )
+    def __init__(self, *, primary_hue=colors.gray, secondary_hue=colors.orange_red,
+                 neutral_hue=colors.slate, text_size=sizes.text_lg,
+                 font=(fonts.GoogleFont("Outfit"), "Arial", "sans-serif"),
+                 font_mono=(fonts.GoogleFont("IBM Plex Mono"), "ui-monospace", "monospace")):
+        super().__init__(primary_hue=primary_hue, secondary_hue=secondary_hue,
+                         neutral_hue=neutral_hue, text_size=text_size,
+                         font=font, font_mono=font_mono)
         super().set(
             background_fill_primary="*primary_50",
             background_fill_primary_dark="*primary_900",
@@ -64,8 +93,7 @@ class OrangeRedTheme(Soft):
             button_primary_background_fill_hover="linear-gradient(90deg, *secondary_600, *secondary_700)",
             button_primary_background_fill_dark="linear-gradient(90deg, *secondary_600, *secondary_700)",
             button_primary_background_fill_hover_dark="linear-gradient(90deg, *secondary_500, *secondary_600)",
-            slider_color="*secondary_500",
-            slider_color_dark="*secondary_600",
+            slider_color="*secondary_500", slider_color_dark="*secondary_600",
             block_title_text_weight="600", block_border_width="3px",
             block_shadow="*shadow_drop_lg", button_primary_shadow="*shadow_drop_lg",
             button_large_padding="11px", color_accent_soft="*primary_100",
@@ -75,40 +103,20 @@ class OrangeRedTheme(Soft):
 orange_red_theme = OrangeRedTheme()
 MAX_SEED = np.iinfo(np.int32).max
 
-# ── Upscaler model registry ───────────────────────────────────────────────────
-# Add new entries here to extend the dropdown — spandrel handles any ESRGAN-
-# family .pth file automatically. scale= is the integer output multiplier.
+# ── Upscaler models, FACE_SWAP_PROMPT, LORA_STYLES — UNCHANGED from previous step
 UPSCALE_MODELS = {
-    "None": {
-        "scale": None, "file": None, "url": None,
-    },
-    "2× — RealESRGAN (balanced)": {
-        "scale": 2,
-        "file": "RealESRGAN_x2plus.pth",
-        "url": "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.1/RealESRGAN_x2plus.pth",
-    },
-    "4× — RealESRGAN (balanced)": {
-        "scale": 4,
-        "file": "RealESRGAN_x4plus.pth",
-        "url": "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth",
-    },
-    "4× — UltraSharp (crisp)": {
-        "scale": 4,
-        "file": "4x-UltraSharpV2.pth",
-        "url": "https://huggingface.co/Kim2091/UltraSharpV2/resolve/main/4x-UltraSharpV2.pth",
-    },
-    "4× — Remacri (natural)": {
-        "scale": 4,
-        "file": "4x_foolhardy_Remacri.pth",
-        "url": "https://huggingface.co/FacehugmanIII/4x_foolhardy_Remacri/resolve/main/4x_foolhardy_Remacri.pth",
-    },
-    "4× — Nomos2 HQ DAT2 (Photography)": {
-        "scale": 4,
-        "file": "4xNomos2_hq_dat2.pth",
-        "url": "https://github.com/Phhofm/models/releases/download/4xNomos2_hq_dat2/4xNomos2_hq_dat2.pth",
-    },    
+    "None": {"scale": None, "file": None, "url": None},
+    "2× — RealESRGAN (balanced)": {"scale": 2, "file": "RealESRGAN_x2plus.pth",
+        "url": "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.1/RealESRGAN_x2plus.pth"},
+    "4× — RealESRGAN (balanced)": {"scale": 4, "file": "RealESRGAN_x4plus.pth",
+        "url": "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth"},
+    "4× — UltraSharp (crisp)": {"scale": 4, "file": "4x-UltraSharpV2.pth",
+        "url": "https://huggingface.co/Kim2091/UltraSharpV2/resolve/main/4x-UltraSharpV2.pth"},
+    "4× — Remacri (natural)": {"scale": 4, "file": "4x_foolhardy_Remacri.pth",
+        "url": "https://huggingface.co/FacehugmanIII/4x_foolhardy_Remacri/resolve/main/4x_foolhardy_Remacri.pth"},
+    "4× — Nomos2 HQ DAT2 (Photography)": {"scale": 4, "file": "4xNomos2_hq_dat2.pth",
+        "url": "https://github.com/Phhofm/models/releases/download/4xNomos2_hq_dat2/4xNomos2_hq_dat2.pth"},
 }
-# ─────────────────────────────────────────────────────────────────────────────
 
 FACE_SWAP_PROMPT = """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.
 FROM PICTURE 1 (strictly preserve):
@@ -122,384 +130,186 @@ FROM PICTURE 2 (strictly preserve identity):
 - Skin: texture, tone, complexion
 The replaced head must seamlessly match Picture 1's lighting and expression while maintaining the complete identity from Picture 2. High quality, photorealistic, sharp details, 4k."""
 
-# MAX number of simultaneous LoRA weight sliders to pre-render in the UI
 MAX_LORA_SLOTS = 6
 
 LORA_STYLES = [
-    {
-        "image": "https://huggingface.co/spaces/prithivMLmods/FLUX.2-Klein-LoRA-Studio/resolve/main/examples/image.webp",
-        "title": "None",
-        "adapter_name": None,
-        "repo": None,
-        "weights": None,
-        "default_prompt": None,
-        "default_weight": 1.0,
-    },
-    {
-        "title": "Klein-Delight-Style",
-        "adapter_name": "klein-delight",
-        "repo": "linoyts/Flux2-Klein-Delight-LoRA",
-        "weights": "pytorch_lora_weights.safetensors",
-        "default_prompt": "Relight the image to remove all existing lighting conditions and replace them with neutral, uniform illumination. Apply soft, evenly distributed lighting with no directional shadows, no harsh highlights, and no dramatic contrast. Maintain the original identity of all subjects exactly—preserve facial structure, skin tone, proportions, expressions, hair, clothing, and textures. Do not alter pose, camera angle, background geometry, or image composition. Lighting should appear balanced, and studio-neutral, similar to diffuse overcast or a soft lightbox setup. Ensure consistent exposure across the entire image with realistic depth and subtle shading only where necessary for form.",
-        "default_weight": 1.0,
-    },
-    {
-        "title": "Klein-Consistency",
-        "adapter_name": "klein-consistency",
-        "repo": "dx8152/Flux2-Klein-9B-Consistency",
-        "weights": "Klein-consistency.safetensors",
-        "default_prompt": None,
-        "default_weight": 0.3,
-    },
-    {
-        "title": "Best-Face-Swap",
-        "adapter_name": "face-swap",
-        "repo": "Alissonerdx/BFS-Best-Face-Swap",
-        "weights": "bfs_head_v1_flux-klein_9b_step3750_rank64.safetensors",
-        "default_prompt": FACE_SWAP_PROMPT,
-        "default_weight": 1.0,
-    },
-    {
-        "title": "NSFW v2",
-        "adapter_name": "nsfw-v2",
-        "repo": "diroverflo/FLux_Klein_9B_NSFW",
-        "weights": "Flux Klein - NSFW v2.safetensors",
-        "default_prompt": None,
-        "default_weight": 1.0,
-    },
-    {
-        "title": "Ultimate Upscaler Klein-9b",
-        "adapter_name": "Ultimate Upscaler",
-        "repo": "loras",
-        "weights": "Flux2-Klein-Image-RestoreV1.safetensors",
-        "default_prompt": "restore the image quality, remove any compression artefacts, remove any haze and soft edges, enrich the original with new intricate detail in all textures and surfaces creating a professional photorealistic photograph with natural lighting and skin texture.",
-        "default_weight": 1.0,
-    },    
-    {
-        "title": "High Resolution",
-        "adapter_name": "High Resolution",
-        "repo": "loras",
-        "weights": "HighResolution9B.safetensors",
-        "default_prompt": "High Resolution",
-        "default_weight": 1.0,
-    },       
-    {
-        "title": "Female Asshole",
-        "adapter_name": "Female Asshole",
-        "repo": "loras",
-        "weights": "femaleasshole-f2-klein-9b.safetensors",
-        "default_prompt": None,
-        "default_weight": 1.0,
-    },        
-    {
-        "title": "Realistic Nudes",
-        "adapter_name": "Realistic Nudes",
-        "repo": "loras",
-        "weights": "realistic_nudes_klein_v3.safetensors",
-        "default_prompt": None,
-        "default_weight": 1.0,
-    },            
-    {
-        "title": "Perky Pointy Puffy Breasts",
-        "adapter_name": "Perky Pointy Puffy Breasts",
-        "repo": "loras",
-        "weights": "PerkyPointyPuffy_v1.1_small_pointy_breasts_large_puffy_nipples.safetensors",
-        "default_prompt": "Small pointy breasts with large puffy nipples",
-        "default_weight": 1.0,
-    },       
-    {
-        "title": "Flat Chested",
-        "adapter_name": "Flat Chested",
-        "repo": "loras",
-        "weights": "Flux2-Klein-9b-FlatChested-v1.safetensors",
-        "default_prompt": "flat chested",
-        "default_weight": 1.5,
-    },           
-    {
-        "title": "Controllight",
-        "adapter_name": "Controllight",
-        "repo": "ControlLight/ControlLight",
-        "weights": "controllight.safetensors",
-        "default_prompt": None,
-        "default_weight": 1.0,
-    },      
-    {
-        "title": "RefControl - Depth",
-        "adapter_name": "RefConDep",
-        "repo": "thedeoxen/refcontrol-FLUX.2-klein-9B-reference-depth-lora",
-        "weights": "flux2_klein_9b_refcontrol_depth.safetensors",
-        "default_prompt": "refcontrol",
-        "default_weight": 1.0,
-    },         
-    {
-        "title": "RefControl - Pose",
-        "adapter_name": "RefConPos",
-        "repo": "thedeoxen/refcontrol-FLUX.2-klein-9B-reference-pose-lora",
-        "weights": "refcontrol_v2_poses.safetensors",
-        "default_prompt": "apply pose from image 1 with reference from image 2",
-        "default_weight": 1.0,
-    },    
+    {"image": "https://huggingface.co/spaces/prithivMLmods/FLUX.2-Klein-LoRA-Studio/resolve/main/examples/image.webp",
+     "title": "None", "adapter_name": None, "repo": None, "weights": None,
+     "default_prompt": None, "default_weight": 1.0},
+    {"title": "Klein-Delight-Style", "adapter_name": "klein-delight",
+     "repo": "linoyts/Flux2-Klein-Delight-LoRA", "weights": "pytorch_lora_weights.safetensors",
+     "default_prompt": "Relight the image to remove all existing lighting conditions and replace them with neutral, uniform illumination. Apply soft, evenly distributed lighting with no directional shadows, no harsh highlights, and no dramatic contrast. Maintain the original identity of all subjects exactly—preserve facial structure, skin tone, proportions, expressions, hair, clothing, and textures. Do not alter pose, camera angle, background geometry, or image composition. Lighting should appear balanced, and studio-neutral, similar to diffuse overcast or a soft lightbox setup. Ensure consistent exposure across the entire image with realistic depth and subtle shading only where necessary for form.",
+     "default_weight": 1.0},
+    {"title": "Klein-Consistency", "adapter_name": "klein-consistency",
+     "repo": "dx8152/Flux2-Klein-9B-Consistency", "weights": "Klein-consistency.safetensors",
+     "default_prompt": None, "default_weight": 0.3},
+    {"title": "Best-Face-Swap", "adapter_name": "face-swap",
+     "repo": "Alissonerdx/BFS-Best-Face-Swap",
+     "weights": "bfs_head_v1_flux-klein_9b_step3750_rank64.safetensors",
+     "default_prompt": FACE_SWAP_PROMPT, "default_weight": 1.0},
+    {"title": "NSFW v2", "adapter_name": "nsfw-v2",
+     "repo": "diroverflo/FLux_Klein_9B_NSFW", "weights": "Flux Klein - NSFW v2.safetensors",
+     "default_prompt": None, "default_weight": 1.0},
+    {"title": "Ultimate Upscaler Klein-9b", "adapter_name": "Ultimate Upscaler",
+     "repo": "loras", "weights": "Flux2-Klein-Image-RestoreV1.safetensors",
+     "default_prompt": "restore the image quality, remove any compression artefacts, remove any haze and soft edges, enrich the original with new intricate detail in all textures and surfaces creating a professional photorealistic photograph with natural lighting and skin texture.",
+     "default_weight": 1.0},
+    {"title": "High Resolution", "adapter_name": "High Resolution",
+     "repo": "loras", "weights": "HighResolution9B.safetensors",
+     "default_prompt": "High Resolution", "default_weight": 1.0},
+    {"title": "Female Asshole", "adapter_name": "Female Asshole",
+     "repo": "loras", "weights": "femaleasshole-f2-klein-9b.safetensors",
+     "default_prompt": None, "default_weight": 1.0},
+    {"title": "Realistic Nudes", "adapter_name": "Realistic Nudes",
+     "repo": "loras", "weights": "realistic_nudes_klein_v3.safetensors",
+     "default_prompt": None, "default_weight": 1.0},
+    {"title": "Perky Pointy Puffy Breasts", "adapter_name": "Perky Pointy Puffy Breasts",
+     "repo": "loras", "weights": "PerkyPointyPuffy_v1.1_small_pointy_breasts_large_puffy_nipples.safetensors",
+     "default_prompt": "Small pointy breasts with large puffy nipples", "default_weight": 1.0},
+    {"title": "Flat Chested", "adapter_name": "Flat Chested",
+     "repo": "loras", "weights": "Flux2-Klein-9b-FlatChested-v1.safetensors",
+     "default_prompt": "flat chested", "default_weight": 1.5},
+    {"title": "Controllight", "adapter_name": "Controllight",
+     "repo": "ControlLight/ControlLight", "weights": "controllight.safetensors",
+     "default_prompt": None, "default_weight": 1.0},
+    {"title": "RefControl - Depth", "adapter_name": "RefConDep",
+     "repo": "thedeoxen/refcontrol-FLUX.2-klein-9B-reference-depth-lora",
+     "weights": "flux2_klein_9b_refcontrol_depth.safetensors",
+     "default_prompt": "refcontrol", "default_weight": 1.0},
+    {"title": "RefControl - Pose", "adapter_name": "RefConPos",
+     "repo": "thedeoxen/refcontrol-FLUX.2-klein-9B-reference-pose-lora",
+     "weights": "refcontrol_v2_poses.safetensors",
+     "default_prompt": "apply pose from image 1 with reference from image 2",
+     "default_weight": 1.0},
 ]
 
-# LOADED_ADAPTERS is the only piece of LoRA state that's legitimately global:
-# it just tracks which adapter names have been loaded onto the shared `pipe`
-# at least once, so we don't re-download/re-attach weights on every call.
-# It is NOT used to decide what a user can select — that comes from each
-# session's own `dynamic_loras_state` (a gr.State dict), so one user's custom
-# HuggingFace LoRA never shows up in another user's checklist.
 LOADED_ADAPTERS = set()
 
 
 def get_all_styles(dynamic_loras):
-    all_styles = list(LORA_STYLES)
-    for dynamic_lora in (dynamic_loras or {}).values():
-        all_styles.append(dynamic_lora)
-    return all_styles
+    return list(LORA_STYLES) + list((dynamic_loras or {}).values())
 
 def get_selectable_styles(dynamic_loras):
-    """Return all styles except 'None' for the checkbox selector."""
     return [s for s in get_all_styles(dynamic_loras) if s["adapter_name"] is not None]
 
 def get_style_by_title(title, dynamic_loras):
-    for style in get_all_styles(dynamic_loras):
-        if style["title"] == title:
-            return style
+    for s in get_all_styles(dynamic_loras):
+        if s["title"] == title:
+            return s
     return None
 
-def get_style_by_adapter_name(adapter_name, dynamic_loras):
-    for style in get_all_styles(dynamic_loras):
-        if style["adapter_name"] == adapter_name:
-            return style
-    return None
 
 print(f"Loading FLUX.2 Klein {MODEL_VARIANT} from {_MODEL_REPO}...")
 pipe = _PipeClass.from_pretrained(_MODEL_REPO, torch_dtype=torch.bfloat16).to(device)
 print(f"Model loaded successfully: FLUX.2 Klein {MODEL_VARIANT}")
 
-def compute_base_dimensions(image):
-    """Scale image so longest side fits within 1024, aligned to 16px grid."""
-    if image is None:
-        return 1024, 1024
-    original_width, original_height = image.size
-    scale = min(1024 / original_width, 1024 / original_height)
-    new_width = (int(original_width * scale) // 16) * 16
-    new_height = (int(original_height * scale) // 16) * 16
-    return new_width, new_height
-
-# Keep old name as alias so nothing inside infer breaks if called directly
-update_dimensions_on_upload = compute_base_dimensions
-
-def fix_orientation(img):
-    """Correct EXIF orientation in-place (returns a new image).
-    Browsers/OS viewers respect the EXIF orientation tag, but PIL does not
-    apply it by default — so an image can *look* correctly oriented
-    everywhere you view it while the actual pixel data fed to the model is
-    rotated 90/180/270°. This normalises the pixels to match what you see."""
-    if img is None:
-        return None
-    return ImageOps.exif_transpose(img)
 
-def on_base_image_change(img):
-    """Update the size/dimension preview when the base image changes."""
-    if img is None:
-        return "*No base image uploaded yet*"
-    try:
-        pil_img = img if isinstance(img, Image.Image) else Image.open(img)
-        orig_w, orig_h = pil_img.size
-        base_w, base_h = compute_base_dimensions(pil_img)
-        return f"Input: **{orig_w} × {orig_h}** px → Output (pre-upscale): **{base_w} × {base_h}** px"
-    except Exception as e:
-        return f"*Could not read dimensions: {e}*"
-
-def on_reference_change(images):
-    """Update the reference image count info when the gallery changes."""
-    if not images:
-        return "📷 No reference images"
-    count = len(images)
-    return f"📷 {count} reference image{'s' if count != 1 else ''} uploaded"
-
-def process_images(base_image, reference_images):
-    """Build the ordered list of PIL images sent to the pipeline:
-    [base_image] + reference_images. EXIF orientation is corrected on every
-    image, regardless of whether it came in via direct upload or the editor."""
-    pil_images = []
-    if base_image is not None:
-        try:
-            img = base_image if isinstance(base_image, Image.Image) else Image.open(base_image)
-            pil_images.append(fix_orientation(img).convert("RGB"))
-        except Exception as e:
-            print(f"Skipping invalid base image: {e}")
-    for item in (reference_images or []):
-        try:
-            path_or_img = item[0] if isinstance(item, (tuple, list)) else item
-            if isinstance(path_or_img, Image.Image):
-                img = path_or_img
-            elif isinstance(path_or_img, str):
-                img = Image.open(path_or_img)
-            else:
-                img = Image.open(path_or_img.name)
-            pil_images.append(fix_orientation(img).convert("RGB"))
-        except Exception as e:
-            print(f"Skipping invalid reference image: {e}")
-    return pil_images
-
-def fix_editor_orientation(editor_value):
-    """Runs on upload into the Image Editor — corrects EXIF orientation on
-    the background layer before the user starts cropping, so what they see
-    in the crop box matches the actual pixel data."""
-    if not editor_value or editor_value.get("background") is None:
-        return editor_value
-    bg = editor_value["background"]
-    if isinstance(bg, Image.Image):
-        editor_value["background"] = fix_orientation(bg)
-    return editor_value
-
-def _editor_composite(editor_value):
-    if not editor_value or editor_value.get("composite") is None:
-        raise gr.Error("Upload and crop an image in the editor first.")
-    composite = editor_value["composite"]
-    if isinstance(composite, np.ndarray):
-        composite = Image.fromarray(composite)
-    return composite.convert("RGB")
-
-def send_editor_to_base(editor_value):
-    composite = _editor_composite(editor_value)
-    gr.Info("Sent to Base Image")
-    return composite
-
-def send_editor_to_reference(editor_value, current_gallery):
-    composite = _editor_composite(editor_value)
-    current = list(current_gallery or [])
-    current.append(composite)
-    gr.Info("Added to Reference Images")
-    return current
-
-# ── LoRA selection → update weight sliders + LoRA prompt display ──────────────
+# ── UI helper callbacks ──────────────────────────────────────────────────────
 
 def update_weight_sliders(selected_titles, dynamic_loras):
-    """
-    Given a list of selected LoRA titles, return updated visibility/label/value
-    for each of the MAX_LORA_SLOTS pre-rendered sliders, plus an update for the
-    LoRA prompt display box (visible when any selected LoRA has a default prompt).
-
-    The user's own prompt is never touched here — LoRA prompts are shown
-    separately and appended at inference time.
-    """
-    selected_styles = [get_style_by_title(t, dynamic_loras) for t in (selected_titles or []) if get_style_by_title(t, dynamic_loras)]
-    slider_updates = []
+    selected = [get_style_by_title(t, dynamic_loras) for t in (selected_titles or [])
+                if get_style_by_title(t, dynamic_loras)]
+    updates = []
     for i in range(MAX_LORA_SLOTS):
-        if i < len(selected_styles):
-            style = selected_styles[i]
-            slider_updates.append(gr.update(
-                visible=True,
-                label=f"{style['title']} — weight",
-                value=style.get("default_weight", 1.0),
-            ))
+        if i < len(selected):
+            s = selected[i]
+            updates.append(gr.update(visible=True,
+                                     label=f"{s['title']} — weight",
+                                     value=s.get("default_weight", 1.0)))
         else:
-            slider_updates.append(gr.update(visible=False, value=1.0))
+            updates.append(gr.update(visible=False, value=1.0))
+    prompts = [s.get("default_prompt") for s in selected if s.get("default_prompt")]
+    if prompts:
+        return updates + [gr.update(value="\n\n".join(prompts), visible=True)]
+    return updates + [gr.update(value="", visible=False)]
 
-    # Collect default prompts from ALL selected LoRAs (not just one)
-    lora_prompts = [s.get("default_prompt") for s in selected_styles if s.get("default_prompt")]
-    if lora_prompts:
-        combined = "\n\n".join(lora_prompts)
-        lora_prompt_update = gr.update(value=combined, visible=True)
-    else:
-        lora_prompt_update = gr.update(value="", visible=False)
 
-    return slider_updates + [lora_prompt_update]
+def on_canvas_mode_change(mode):
+    """Custom W/H sliders only relevant when mode == Custom."""
+    is_custom = (mode == "Custom")
+    return gr.update(visible=is_custom), gr.update(visible=is_custom)
 
-# ── Inference ─────────────────────────────────────────────────────────────────
 
-def _upscale_tiled(model_fn, img_t: torch.Tensor, tile: int = 512, overlap: int = 32) -> torch.Tensor:
-    """
-    Run the upscaler model patch-by-patch to cap peak VRAM usage.
-    Overlapping tiles are averaged so there are no seam artifacts.
-    Output is assembled on CPU so only one tile lives on GPU at a time.
-    """
-    _, c, h, w = img_t.shape
+def on_fit_mode_change(fit_mode):
+    """Pad colour swatch only relevant for Pad (color)."""
+    return gr.update(visible=(fit_mode == "Pad (color)"))
 
-    # Probe output scale with a tiny crop rather than hard-coding it
+
+def on_batch_vary_change(vary_mode):
+    """Sweep range only relevant for the LoRA sweep mode."""
+    is_sweep = (vary_mode == "Sweep first LoRA weight")
+    return gr.update(visible=is_sweep), gr.update(visible=is_sweep)
+
+
+def on_gallery_select(evt: gr.SelectData, gallery_value):
+    """Remember which gallery item is selected so Send→* uses it."""
+    if evt is None or gallery_value is None or evt.index is None:
+        return None
+    try:
+        item = gallery_value[evt.index]
+    except (IndexError, TypeError):
+        return None
+    return item[0] if isinstance(item, (list, tuple)) else item
+
+
+# ── Upscaler (tiled) — unchanged ─────────────────────────────────────────────
+
+def _upscale_tiled(model_fn, img_t, tile=512, overlap=32):
+    _, c, h, w = img_t.shape
     with torch.no_grad():
         probe = model_fn(img_t[:, :, :min(4, h), :min(4, w)])
     scale = probe.shape[-1] // min(4, w)
-    del probe
-    torch.cuda.empty_cache()
-
+    del probe; torch.cuda.empty_cache()
     out_h, out_w = h * scale, w * scale
     canvas  = torch.zeros(1, c, out_h, out_w, dtype=torch.float32)
     weights = torch.zeros(1, 1, out_h, out_w, dtype=torch.float32)
-
     step = max(tile - overlap, 1)
-    # Ensure the last tile always reaches the edge
-    ys = list(range(0, max(h - tile, 0), step)) + [max(h - tile, 0)]
-    xs = list(range(0, max(w - tile, 0), step)) + [max(w - tile, 0)]
-    ys = sorted(set(ys))
-    xs = sorted(set(xs))
-
+    ys = sorted(set(list(range(0, max(h - tile, 0), step)) + [max(h - tile, 0)]))
+    xs = sorted(set(list(range(0, max(w - tile, 0), step)) + [max(w - tile, 0)]))
     for y0 in ys:
         for x0 in xs:
-            y1 = min(y0 + tile, h)
-            x1 = min(x0 + tile, w)
-            patch = img_t[:, :, y0:y1, x0:x1]
+            y1 = min(y0 + tile, h); x1 = min(x0 + tile, w)
             with torch.no_grad():
-                out = model_fn(patch).cpu().float()
+                out = model_fn(img_t[:, :, y0:y1, x0:x1]).cpu().float()
             oy0, ox0 = y0 * scale, x0 * scale
             oy1, ox1 = y1 * scale, x1 * scale
             canvas[:, :, oy0:oy1, ox0:ox1]  += out
             weights[:, :, oy0:oy1, ox0:ox1] += 1.0
-
     return (canvas / weights.clamp(min=1)).clamp(0, 1)
 
 
-def apply_realesrgan(image: Image.Image, model_key: str) -> Image.Image:
-    """
-    Upscale a PIL image using the model selected in UPSCALE_MODELS.
-    Weights are downloaded once to /tmp/realesrgan_weights/ and reused.
-    """
+def apply_realesrgan(image, model_key):
     cfg = UPSCALE_MODELS[model_key]
-
     try:
         from spandrel import ImageModelDescriptor, ModelLoader
     except ImportError:
         raise gr.Error("spandrel is not installed. Add 'spandrel' to requirements.txt.")
-
     import urllib.request
-
     cache_dir = "/tmp/realesrgan_weights"
     os.makedirs(cache_dir, exist_ok=True)
     cache_path = os.path.join(cache_dir, cfg["file"])
-
     if not os.path.exists(cache_path):
         print(f"Downloading {cfg['file']}…")
         urllib.request.urlretrieve(cfg["url"], cache_path)
-
     model = ModelLoader().load_from_file(cache_path)
     if not isinstance(model, ImageModelDescriptor):
         raise gr.Error(f"Loaded model is not a single-image descriptor: {type(model)}")
-
     sr_model = model.model.to(device).eval()
-
     img_np = np.array(image).astype(np.float32) / 255.0
     img_t  = torch.from_numpy(img_np).permute(2, 0, 1).unsqueeze(0).to(device)
-
     out_t  = _upscale_tiled(sr_model, img_t, tile=512, overlap=32)
-
-    # Free upscaler weights immediately — FLUX pipeline stays resident
     del sr_model, img_t
-    gc.collect()
-    torch.cuda.empty_cache()
-
+    gc.collect(); torch.cuda.empty_cache()
     out_np = out_t.squeeze(0).permute(1, 2, 0).numpy()
     return Image.fromarray((out_np * 255).astype(np.uint8))
 
 
-# ── Logging helper — fires a daemon thread so logging never blocks the user ──
+# ── Logging ──────────────────────────────────────────────────────────────────
 
 def _spawn_log(pil_images, result_image, prompt, seed, steps, guidance_scale,
                width, height, duration, success, error="",
                lora_titles=None, lora_weights=None, upscale_factor="None",
                lora_prompt_text=""):
-    
     if os.environ.get("ENABLE_LOGGING", "No").strip().lower() != "yes":
         return
     threading.Thread(
@@ -512,186 +322,283 @@ def _spawn_log(pil_images, result_image, prompt, seed, steps, guidance_scale,
     ).start()
 
 
-def infer(
-    base_image,
-    reference_images,
-    prompt,
-    lora_prompt_text,
-    custom_prompt_text,
-    selected_titles,
-    seed,
-    randomize_seed,
-    guidance_scale,
-    steps,
-    upscale_factor,
-    dynamic_loras,
-    *slider_values,
-    progress=gr.Progress(track_tqdm=True)
-):
-    """CPU-side wrapper — handles pre/post processing and logging.
-    The actual GPU work is delegated to _infer_gpu so that the
-    _spawn_log daemon thread survives after the @spaces.GPU subprocess exits."""
-    gc.collect()
-    torch.cuda.empty_cache()
-
-    # Guard against stale MCP schema sending the old output_scale float value
-    if not isinstance(upscale_factor, str) or upscale_factor not in UPSCALE_MODELS:
-        upscale_factor = "None"
-
-    if base_image is None:
-        raise gr.Error("Please upload a base image.")
-
-    pil_images = process_images(base_image, reference_images)
-    if not pil_images:
-        raise gr.Error("Could not process uploaded images.")
-
-    selected_titles = selected_titles or []
-
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
-
-    # Collect LoRA metadata for logging before entering GPU
-    active_styles = [get_style_by_title(t, dynamic_loras) for t in selected_titles
-                     if get_style_by_title(t, dynamic_loras) and get_style_by_title(t, dynamic_loras)["adapter_name"] is not None]
-    log_lora_titles = [s["title"] for s in active_styles]
-    log_lora_weights = [float(w) for w in slider_values[:len(active_styles)]]
-    log_lora_prompt = (lora_prompt_text or "").strip()
-
-    t0 = time.perf_counter()
-    try:
-        image, seed, width, height = _infer_gpu(
-            pil_images, prompt, lora_prompt_text, custom_prompt_text, selected_titles,
-            seed, guidance_scale, steps, upscale_factor, dynamic_loras,
-            *slider_values, progress=progress
-        )
-        duration = time.perf_counter() - t0
-        # _spawn_log runs on CPU main process so the daemon thread survives
-        _spawn_log(pil_images, image, prompt, seed, steps, guidance_scale,
-                   width, height, duration, True,
-                   lora_titles=log_lora_titles, lora_weights=log_lora_weights,
-                   upscale_factor=upscale_factor, lora_prompt_text=log_lora_prompt)
-        return image, str(seed)
-    except Exception as e:
-        duration = time.perf_counter() - t0
-        _spawn_log(pil_images, None, prompt, seed, steps, guidance_scale,
-                   0, 0, duration, False, str(e),
-                   lora_titles=log_lora_titles, lora_weights=log_lora_weights,
-                   upscale_factor=upscale_factor, lora_prompt_text=log_lora_prompt)
-        raise
-
+# ── GPU step (shared by single, batch, and bulk) ─────────────────────────────
 
 @spaces.GPU
 def _infer_gpu(
-    pil_images,
-    prompt,
-    lora_prompt_text,
-    custom_prompt_text,
-    selected_titles,
-    seed,
-    guidance_scale,
-    steps,
-    upscale_factor,
-    dynamic_loras,
-    *slider_values,
-    progress=gr.Progress(track_tqdm=True)
+    pil_images, prompt, lora_prompt_text, custom_prompt_text, selected_titles,
+    seed, guidance_scale, steps, upscale_factor,
+    canvas_mode, custom_width, custom_height,
+    canvas_fit_mode, pad_color, dynamic_loras,
+    *slider_values, progress=gr.Progress(track_tqdm=True),
 ):
-    """GPU-side inference — all LoRA loading, pipe() calls, and upscaling happen here."""
-
-    # Face-swap validation
     if "Best-Face-Swap" in selected_titles:
         if len(pil_images) < 2:
-            raise gr.Error("Face Swap requires 2 images: a Base image and one Reference image (used as the face source).")
+            raise gr.Error("Face Swap requires 2 images: a Base image and one Reference image.")
         if len(pil_images) > 2:
             gr.Warning("Face Swap uses only the Base image and the first Reference image.")
             pil_images = pil_images[:2]
 
-    # Build active adapter list (skip None)
-    active_styles = [get_style_by_title(t, dynamic_loras) for t in selected_titles if get_style_by_title(t, dynamic_loras) and get_style_by_title(t, dynamic_loras)["adapter_name"] is not None]
+    active_styles = [get_style_by_title(t, dynamic_loras) for t in selected_titles
+                     if get_style_by_title(t, dynamic_loras)
+                     and get_style_by_title(t, dynamic_loras)["adapter_name"] is not None]
     weights = list(slider_values[:len(active_styles)])
 
     if not active_styles:
-        print("No LoRA selected — running base model.")
         pipe.disable_lora()
     else:
-        for i, style in enumerate(active_styles):
-            adapter_name = style["adapter_name"]
-            if adapter_name not in LOADED_ADAPTERS:
-                print(f"Loading adapter: {style['title']}")
+        for style in active_styles:
+            an = style["adapter_name"]
+            if an not in LOADED_ADAPTERS:
                 try:
-                    pipe.load_lora_weights(
-                        style["repo"],
-                        weight_name=style["weights"],
-                        adapter_name=adapter_name,
-                    )
-                    LOADED_ADAPTERS.add(adapter_name)
+                    pipe.load_lora_weights(style["repo"], weight_name=style["weights"], adapter_name=an)
+                    LOADED_ADAPTERS.add(an)
                 except Exception as e:
                     raise gr.Error(f"Failed to load {style['title']}: {e}")
+        pipe.set_adapters([s["adapter_name"] for s in active_styles],
+                          adapter_weights=[float(w) for w in weights])
 
-        adapter_names = [s["adapter_name"] for s in active_styles]
-        adapter_weights = [float(w) for w in weights]
-        print(f"Activating adapters: {list(zip(adapter_names, adapter_weights))}")
-        pipe.set_adapters(adapter_names, adapter_weights=adapter_weights)
-
-
-    # Combine user prompt with LoRA default prompts and any selected custom prompts
-    user_prompt = (prompt or "").strip()
-    lora_extra = (lora_prompt_text or "").strip()
-    custom_extra = (custom_prompt_text or "").strip()
-    parts = [p for p in [user_prompt, lora_extra, custom_extra] if p]
-    full_prompt = "\n".join(parts)
+    full_prompt = "\n".join(p for p in [
+        (prompt or "").strip(),
+        (lora_prompt_text or "").strip(),
+        (custom_prompt_text or "").strip(),
+    ] if p)
 
-    # Always generate at base 1024-capped resolution
-    width, height = compute_base_dimensions(pil_images[0])
-    print(f"Generating at: {width} × {height} px")
+    width, height = compute_canvas_dimensions(pil_images[0], canvas_mode, custom_width, custom_height)
+    print(f"Generating at: {width}×{height} (canvas={canvas_mode}, fit={canvas_fit_mode})")
 
-    processed_images = [img.resize((width, height), Image.LANCZOS).convert("RGB") for img in pil_images]
-    image_input = processed_images if len(processed_images) > 1 else processed_images[0]
+    processed = [fit_to_canvas(img, width, height, canvas_fit_mode, pad_color) for img in pil_images]
+    image_input = processed if len(processed) > 1 else processed[0]
 
     try:
-        # Flux2KleinKVPipeline is a distilled model and does not accept
-        # guidance_scale — only pass it for the non-KV 9B variant.
-        call_kwargs = dict(
-            image=image_input,
-            prompt=full_prompt,
-            width=width,
-            height=height,
-            num_inference_steps=steps,
-            generator=torch.Generator(device=device).manual_seed(seed),
-        )
+        kwargs = dict(image=image_input, prompt=full_prompt,
+                      width=width, height=height,
+                      num_inference_steps=steps,
+                      generator=torch.Generator(device=device).manual_seed(seed))
         if MODEL_VARIANT != "9B-KV":
-            call_kwargs["guidance_scale"] = guidance_scale
-        image = pipe(**call_kwargs).images[0]
+            kwargs["guidance_scale"] = guidance_scale
+        image = pipe(**kwargs).images[0]
     except Exception as e:
         raise gr.Error(f"Inference failed: {e}")
 
-    # ── Post-generation upscaling ─────────────────────────────────────────────
     if upscale_factor and upscale_factor != "None":
-        scale_int = UPSCALE_MODELS[upscale_factor]["scale"]
-        print(f"Upscaling {scale_int}× with {upscale_factor}…")
-        # Flush FLUX pipeline residuals before loading upscaler weights
-        gc.collect()
-        torch.cuda.synchronize()
-        torch.cuda.empty_cache()
+        gc.collect(); torch.cuda.synchronize(); torch.cuda.empty_cache()
         try:
             image = apply_realesrgan(image, upscale_factor)
-            print(f"Upscaled to: {image.width} × {image.height} px")
         except Exception as e:
-            gr.Warning(f"Upscaling failed, returning 1024px result: {e}")
-    # ──────────────────────────────────────���──────────────────────────────────
+            gr.Warning(f"Upscaling failed, returning {width}×{height} result: {e}")
 
-    gc.collect()
-    torch.cuda.empty_cache()
+    gc.collect(); torch.cuda.empty_cache()
     return image, seed, width, height
 
-# ── Dynamic LoRA loader ───────────────────────────────────────────────────────
+
+def _meta_for(prompt, seed, steps, guidance_scale, width, height, upscale_factor,
+              canvas_mode, canvas_fit_mode, lora_titles, lora_weights,
+              lora_prompt, custom_prompt, extra=None):
+    meta = {
+        "prompt": (prompt or "").strip(),
+        "seed": seed, "steps": steps,
+        "guidance_scale": guidance_scale if MODEL_VARIANT != "9B-KV" else None,
+        "width": width, "height": height,
+        "model": f"FLUX.2-Klein-{MODEL_VARIANT}",
+        "upscale_factor": upscale_factor,
+        "canvas_mode": canvas_mode, "canvas_fit_mode": canvas_fit_mode,
+        "loras": list(zip(lora_titles or [], lora_weights or [])),
+        "lora_prompt": (lora_prompt or "").strip(),
+        "custom_prompt": (custom_prompt or "").strip(),
+    }
+    if extra:
+        meta.update(extra)
+    return meta
+
+
+# ── Single / batch infer (generator → streams into gr.Gallery) ───────────────
+
+def infer(
+    base_image, reference_images, prompt, lora_prompt_text, custom_prompt_text,
+    selected_titles, seed, randomize_seed, guidance_scale, steps, upscale_factor,
+    canvas_mode, custom_width, custom_height, canvas_fit_mode, pad_color,
+    batch_count, batch_vary, sweep_min, sweep_max,
+    dynamic_loras, *slider_values, progress=gr.Progress(track_tqdm=True),
+):
+    """Generator. Streams a list of PNG paths into the output gallery, one
+    result at a time. Each iteration is its own @spaces.GPU call, so a
+    ZeroGPU quota wall mid-batch only loses the in-progress item — earlier
+    PNGs are already saved to disk and surfaced in the gallery."""
+    gc.collect(); torch.cuda.empty_cache()
+    if not isinstance(upscale_factor, str) or upscale_factor not in UPSCALE_MODELS:
+        upscale_factor = "None"
+    if base_image is None:
+        raise gr.Error("Please upload a base image.")
+    pil_images = process_images(base_image, reference_images)
+    if not pil_images:
+        raise gr.Error("Could not process uploaded images.")
+
+    selected_titles = selected_titles or []
+    batch_count = max(1, int(batch_count))
+
+    # Pre-plan per-iteration seed and weight overrides — keeps the loop simple
+    # and lets us put the LoRA-sweep values into PNG metadata cleanly.
+    base_seed = random.randint(0, MAX_SEED) if randomize_seed else int(seed)
+    seeds, weight_overrides = [], []
+    for i in range(batch_count):
+        if batch_vary == "Sequential seed (+1 each)":
+            seeds.append((base_seed + i) % (MAX_SEED + 1)); weight_overrides.append(None)
+        elif batch_vary == "Sweep first LoRA weight":
+            seeds.append(base_seed)
+            t = i / max(batch_count - 1, 1)
+            weight_overrides.append((0, float(sweep_min) + t * (float(sweep_max) - float(sweep_min))))
+        else:  # "Random seed each run" (default)
+            seeds.append(random.randint(0, MAX_SEED)); weight_overrides.append(None)
+
+    active_styles = [get_style_by_title(t, dynamic_loras) for t in selected_titles
+                     if get_style_by_title(t, dynamic_loras)
+                     and get_style_by_title(t, dynamic_loras)["adapter_name"] is not None]
+    log_titles = [s["title"] for s in active_styles]
+
+    results = []
+    last_seed_text = ""
+    for i in range(batch_count):
+        sliders = list(slider_values)
+        if weight_overrides[i] is not None:
+            slot, val = weight_overrides[i]
+            if slot < len(sliders):
+                sliders[slot] = val
+
+        cur_seed = seeds[i]
+        log_weights = [float(w) for w in sliders[:len(active_styles)]]
+        t0 = time.perf_counter()
+        try:
+            image, used_seed, w, h = _infer_gpu(
+                pil_images, prompt, lora_prompt_text, custom_prompt_text, selected_titles,
+                cur_seed, guidance_scale, steps, upscale_factor,
+                canvas_mode, custom_width, custom_height,
+                canvas_fit_mode, pad_color, dynamic_loras,
+                *sliders, progress=progress,
+            )
+            meta = _meta_for(prompt, used_seed, steps, guidance_scale, w, h, upscale_factor,
+                             canvas_mode, canvas_fit_mode, log_titles, log_weights,
+                             lora_prompt_text, custom_prompt_text,
+                             extra={"batch_index": i, "batch_total": batch_count,
+                                    "batch_vary": batch_vary})
+            results.append(save_with_metadata(image, meta))
+            last_seed_text = str(used_seed)
+            _spawn_log(pil_images, image, prompt, used_seed, steps, guidance_scale,
+                       w, h, time.perf_counter() - t0, True,
+                       lora_titles=log_titles, lora_weights=log_weights,
+                       upscale_factor=upscale_factor, lora_prompt_text=lora_prompt_text or "")
+            yield results, last_seed_text
+        except Exception as e:
+            _spawn_log(pil_images, None, prompt, cur_seed, steps, guidance_scale,
+                       0, 0, time.perf_counter() - t0, False, str(e),
+                       lora_titles=log_titles, lora_weights=log_weights,
+                       upscale_factor=upscale_factor, lora_prompt_text=lora_prompt_text or "")
+            yield results, f"Batch {i+1}/{batch_count} failed: {e}"
+            # keep going — earlier results are preserved in the gallery
+
+
+# ── Bulk processing (one input image per iteration) ─────────────────────────
+
+def _new_bulk_workdir() -> str:
+    sid = uuid.uuid4().hex[:8]
+    path = f"/tmp/bulk_{sid}"
+    os.makedirs(path, exist_ok=True)
+    return path
+
+
+def bulk_infer(
+    input_files,
+    prompt, lora_prompt_text, custom_prompt_text, selected_titles,
+    seed, randomize_seed, guidance_scale, steps, upscale_factor,
+    canvas_mode, custom_width, custom_height, canvas_fit_mode, pad_color,
+    dynamic_loras, *slider_values, progress=gr.Progress(),
+):
+    """Process each uploaded image as its own GPU call, streaming results
+    into the bulk output gallery as soon as they complete. Outputs and a CSV
+    manifest are written to /tmp/bulk_<sid>/ with stable filenames, so a
+    ZeroGPU quota wall mid-run still leaves earlier outputs grabbable from
+    the gallery and from disk."""
+    if not input_files:
+        raise gr.Error("Upload at least one image first.")
+
+    work_dir = _new_bulk_workdir()
+    manifest_path = os.path.join(work_dir, "manifest.csv")
+    with open(manifest_path, "w", newline="") as f:
+        csv.writer(f).writerow([
+            "idx", "input_filename", "output_path", "seed",
+            "width", "height", "success", "error", "duration_sec",
+        ])
+
+    active_styles = [get_style_by_title(t, dynamic_loras) for t in (selected_titles or [])
+                     if get_style_by_title(t, dynamic_loras)
+                     and get_style_by_title(t, dynamic_loras)["adapter_name"] is not None]
+    log_titles = [s["title"] for s in active_styles]
+    log_weights = [float(w) for w in slider_values[:len(active_styles)]]
+
+    results, succeeded, failed = [], 0, 0
+    total = len(input_files)
+
+    for i, path in enumerate(input_files):
+        progress(i / total, desc=f"Image {i+1}/{total}")
+        fname = os.path.basename(path) if isinstance(path, str) else f"input_{i}"
+        t0 = time.perf_counter()
+        try:
+            img = fix_orientation(Image.open(path)).convert("RGB")
+            cur_seed = random.randint(0, MAX_SEED) if randomize_seed else int(seed)
+
+            image, used_seed, w, h = _infer_gpu(
+                [img], prompt, lora_prompt_text, custom_prompt_text, selected_titles or [],
+                cur_seed, guidance_scale, steps, upscale_factor,
+                canvas_mode, custom_width, custom_height,
+                canvas_fit_mode, pad_color, dynamic_loras,
+                *slider_values, progress=progress,
+            )
+
+            # Stable, predictable filename inside the session work_dir so the
+            # user can also find outputs on disk if the UI drops.
+            stem = os.path.splitext(fname)[0]
+            out_path = os.path.join(work_dir, f"{i:03d}_{stem}.png")
+            meta = _meta_for(prompt, used_seed, steps, guidance_scale, w, h, upscale_factor,
+                             canvas_mode, canvas_fit_mode, log_titles, log_weights,
+                             lora_prompt_text, custom_prompt_text,
+                             extra={"bulk_input": fname,
+                                    "bulk_index": i, "bulk_total": total})
+            image.save(out_path, format="PNG", pnginfo=build_pnginfo(meta))
+            results.append(out_path)
+            succeeded += 1
+            duration = time.perf_counter() - t0
+
+            with open(manifest_path, "a", newline="") as f:
+                csv.writer(f).writerow([i, fname, out_path, used_seed, w, h, True, "", f"{duration:.2f}"])
+
+            _spawn_log([img], image, prompt, used_seed, steps, guidance_scale,
+                       w, h, duration, True,
+                       lora_titles=log_titles, lora_weights=log_weights,
+                       upscale_factor=upscale_factor, lora_prompt_text=lora_prompt_text or "")
+
+            status = f"✅ {succeeded}/{total} done ({failed} failed)"
+            yield results, status, gr.update()   # zip not ready yet
+        except Exception as e:
+            failed += 1
+            duration = time.perf_counter() - t0
+            with open(manifest_path, "a", newline="") as f:
+                csv.writer(f).writerow([i, fname, "", "", "", "", False, str(e)[:300], f"{duration:.2f}"])
+            yield results, f"⚠️ Image {i+1} failed: {e} | {succeeded} ok, {failed} failed", gr.update()
+            continue
+
+    # Final pass: bundle a zip. ZIP_STORED because PNGs are already compressed.
+    zip_path = os.path.join(work_dir, "outputs.zip")
+    with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_STORED) as zf:
+        for p in results:
+            zf.write(p, arcname=os.path.basename(p))
+        zf.write(manifest_path, arcname="manifest.csv")
+
+    yield results, f"🎉 Done — {succeeded}/{total} succeeded, {failed} failed", zip_path
+
+
+# ── Dynamic LoRA loader (unchanged from previous step) ───────────────────────
 
 def add_custom_lora(repo_id, weight_name, adapter_name, dynamic_loras_state):
-    """
-    Adds a custom LoRA to *this session's* state only — other users never see
-    it in their selector. The adapter_name attached to the shared pipe is
-    still made globally unique (via a short random suffix) since the pipe
-    object itself is one shared resource across all sessions.
-    """
     dynamic_loras = dict(dynamic_loras_state or {})
     if not repo_id or not repo_id.strip():
         return "Please enter a valid HuggingFace repo ID.", gr.update(), dynamic_loras
@@ -709,74 +616,51 @@ def add_custom_lora(repo_id, weight_name, adapter_name, dynamic_loras_state):
         if not actual_weight:
             available = [f.filename for f in info.siblings if f.filename.endswith(('.safetensors', '.bin'))]
             return f"No weight found. Available: {', '.join(available) or 'None'}", gr.update(), dynamic_loras
-
         base_name = "".join(c if c.isalnum() or c in "-_" else "_" for c in requested_name) if requested_name else "custom"
-
         static_names = {s["adapter_name"] for s in LORA_STYLES if s["adapter_name"]}
-        final_adapter_name = f"{base_name}_{uuid.uuid4().hex[:6]}"
-        while final_adapter_name in static_names or final_adapter_name in LOADED_ADAPTERS:
-            final_adapter_name = f"{base_name}_{uuid.uuid4().hex[:6]}"
-
-        custom_style = {
+        final = f"{base_name}_{uuid.uuid4().hex[:6]}"
+        while final in static_names or final in LOADED_ADAPTERS:
+            final = f"{base_name}_{uuid.uuid4().hex[:6]}"
+        dynamic_loras[final] = {
             "image": "https://huggingface.co/spaces/prithivMLmods/FLUX.2-Klein-LoRA-Studio/resolve/main/examples/image.webp",
-            "title": f"Custom: {base_name}",
-            "adapter_name": final_adapter_name,
-            "repo": repo_id,
-            "weights": actual_weight,
-            "default_prompt": None,
-            "default_weight": 1.0,
+            "title": f"Custom: {base_name}", "adapter_name": final,
+            "repo": repo_id, "weights": actual_weight,
+            "default_prompt": None, "default_weight": 1.0,
         }
-        dynamic_loras[final_adapter_name] = custom_style
         new_choices = [s["title"] for s in get_selectable_styles(dynamic_loras)]
         return f"✅ Added: {base_name} from {repo_id}", gr.update(choices=new_choices), dynamic_loras
     except Exception as e:
         return f"❌ Failed: {str(e)}", gr.update(), dynamic_loras
 
-# ── Custom Prompt manager ─────────────────────────────────────────────────────
+
+# ── Custom prompt manager (unchanged) ────────────────────────────────────────
 
 def add_custom_prompt(name, text, prompts_state, counter_state):
-    """
-    All state is passed in and returned — never touches module-level globals.
-    Each user session gets its own isolated copy via gr.State.
-    """
-    prompts = dict(prompts_state)   # never mutate the state object directly
-    counter = int(counter_state)
+    prompts = dict(prompts_state); counter = int(counter_state)
     text = text.strip() if text else ""
     name = name.strip() if name else ""
     if not text:
         return "Please enter some prompt text.", prompts, counter, gr.update(), gr.update(), gr.update()
     if not name:
-        counter += 1
-        name = f"Prompt {counter}"
+        counter += 1; name = f"Prompt {counter}"
     if name in prompts:
-        return f"⚠️ '{name}' already exists — choose a different name.", prompts, counter, gr.update(), gr.update(), gr.update()
+        return f"⚠️ '{name}' already exists.", prompts, counter, gr.update(), gr.update(), gr.update()
     prompts[name] = text
     choices = list(prompts.keys())
-    return (
-        f"✅ Saved: '{name}'",
-        prompts,                                # updated prompts state
-        counter,                                # updated counter state
-        gr.update(choices=choices),             # custom_prompt_selector
-        gr.update(choices=choices),             # delete_prompt_name dropdown
-        gr.update(value="", interactive=True),  # clear the name field
-    )
+    return (f"✅ Saved: '{name}'", prompts, counter,
+            gr.update(choices=choices), gr.update(choices=choices),
+            gr.update(value="", interactive=True))
 
 def delete_custom_prompt(name, currently_selected, prompts_state):
     prompts = dict(prompts_state)
+    msg = f"🗑️ Deleted: '{name}'" if name and name in prompts else "Nothing to delete."
     if name and name in prompts:
         del prompts[name]
-        msg = f"🗑️ Deleted: '{name}'"
-    else:
-        msg = "Nothing to delete."
     choices = list(prompts.keys())
-    # Drop deleted name from active selection if it was ticked
-    new_selection = [n for n in (currently_selected or []) if n in prompts]
-    return (
-        msg,
-        prompts,                                          # updated prompts state
-        gr.update(choices=choices, value=new_selection),  # custom_prompt_selector
-        gr.update(choices=choices, value=None),           # delete_prompt_name
-    )
+    new_sel = [n for n in (currently_selected or []) if n in prompts]
+    return (msg, prompts,
+            gr.update(choices=choices, value=new_sel),
+            gr.update(choices=choices, value=None))
 
 def update_custom_prompt_display(selected_names, prompts_state):
     if not selected_names:
@@ -787,7 +671,7 @@ def update_custom_prompt_display(selected_names, prompts_state):
     return gr.update(value="", visible=False)
 
 
-# ── UI ────────────────────────────────────────────────────────────────────────
+# ── UI ───────────────────────────────────────────────────────────────────────
 
 css = """
 #col-container { margin: 0 auto; max-width: 980px; }
@@ -796,11 +680,12 @@ css = """
 .lora-weight-row { background: var(--block-background-fill); border-radius: 8px; padding: 4px 12px; margin-bottom: 4px; }
 """
 
-with gr.Blocks(css=css, theme=orange_red_theme) as demo:
-    # Per-session state — each user gets their own isolated copy, never shared across sessions
+# Gradio 6.0: theme/css go on launch(), not Blocks()
+with gr.Blocks() as demo:
     custom_prompts_state = gr.State({})
     custom_prompt_counter_state = gr.State(0)
-    dynamic_loras_state = gr.State({})  # this session's custom HF LoRAs — never shared with other users
+    dynamic_loras_state = gr.State({})
+    selected_output_state = gr.State(None)  # currently-selected gallery item path
 
     with gr.Column(elem_id="col-container"):
         _logging_on = os.environ.get("ENABLE_LOGGING", "No").strip().lower() == "yes"
@@ -809,206 +694,286 @@ with gr.Blocks(css=css, theme=orange_red_theme) as demo:
         gr.Markdown("# **FLUX.2-Klein-LoRA-Studio**", elem_id="main-title")
         gr.Markdown(
             f"Apply one or more [LoRA](https://huggingface.co/models?other=base_model:adapter:black-forest-labs/FLUX.2-klein-9B) "
-            f"adapters to your images using "
-            f"[FLUX.2-Klein-{MODEL_VARIANT}]({_MODEL_REPO}).\n\n"
+            f"adapters using [FLUX.2-Klein-{MODEL_VARIANT}]({_MODEL_REPO}). "
             f"**Model:** `{MODEL_VARIANT}` · **Logging:** {_logging_badge}"
         )
-        gr.Markdown(
-            "*Results may vary when combining multiple LoRAs. "
-            "Uploaded images and metadata may be temporarily stored to help optimise the space — "
-            "stored data is automatically deleted. "
-            "Always comply with HF and model Terms of Service.*"
-        )
 
         with gr.Tabs() as main_tabs:
-            # ── Generate tab ────────────────────────────────────────────────────
+            # ── Generate tab ─────────────────────────────────────────────────
             with gr.Tab("🎨 Generate", id="tab_generate"):
                 with gr.Row(equal_height=False):
-                    # ── Left column ──────────────────────────────────────────────
                     with gr.Column(scale=1):
                         base_image = gr.Image(
-                            label="Base Image",
-                            type="pil",
-                            sources=["upload", "clipboard"],
-                            height=260,
-                            elem_id="base_image",
+                            label="Base Image", type="pil",
+                            sources=["upload", "clipboard"], height=260, elem_id="base_image",
                         )
                         size_info = gr.Markdown("*No image uploaded yet*")
 
                         reference_images = gr.Gallery(
-                            label="Reference Image(s) — optional",
-                            type="filepath",
-                            columns=2, rows=1, height=140,
-                            allow_preview=True,
+                            label="Reference Image(s) — optional", type="filepath",
+                            columns=2, rows=1, height=140, allow_preview=True,
                             elem_id="reference_gallery",
                         )
                         reference_info = gr.Markdown("📷 No reference images")
-                        gr.Markdown(
-                            "*For Face Swap, the first Reference image is used as the face source.*"
-                        )
-
-                        prompt = gr.Text(
-                            label="Prompt",
-                            max_lines=3,
-                            placeholder="Describe the edit, or leave blank for style-only LoRAs",
-                        )
-
-                        # Read-only display for LoRA default prompts — shown when any
-                        # selected LoRA has a default prompt. Combined with user prompt
-                        # at inference time; never overwrites what the user typed.
-                        lora_prompt_display = gr.Textbox(
-                            label="LoRA Default Prompts (auto-appended)",
-                            info="These prompts come from your selected LoRAs and are appended to your prompt above during generation.",
-                            interactive=False,
-                            visible=False,
-                            lines=3,
-                        )
-
-                        # Read-only display for active custom prompts — shown when any
-                        # saved custom prompt is ticked in the Custom Prompts accordion.
-                        custom_prompt_display = gr.Textbox(
-                            label="Custom Prompts (auto-appended)",
-                            info="These are your selected saved prompts, appended alongside LoRA prompts during generation.",
-                            interactive=False,
-                            visible=False,
-                            lines=3,
-                        )
-
+                        gr.Markdown("*For Face Swap, the first Reference image is used as the face source.*")
+
+                        prompt = gr.Text(label="Prompt", max_lines=3,
+                                         placeholder="Describe the edit, or leave blank for style-only LoRAs")
+                        lora_prompt_display = gr.Textbox(label="LoRA Default Prompts (auto-appended)",
+                                                         interactive=False, visible=False, lines=3)
+                        custom_prompt_display = gr.Textbox(label="Custom Prompts (auto-appended)",
+                                                           interactive=False, visible=False, lines=3)
                         run_button = gr.Button("▶ Generate", variant="primary", size="lg")
 
                         with gr.Accordion("⚙️ Advanced Settings", open=False):
                             seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0)
                             randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
-                            guidance_scale = gr.Slider(label="Guidance Scale", minimum=0.0, maximum=10.0, step=0.1, value=1.0, visible=MODEL_VARIANT != "9B-KV")
+                            guidance_scale = gr.Slider(label="Guidance Scale", minimum=0.0, maximum=10.0,
+                                                       step=0.1, value=1.0, visible=MODEL_VARIANT != "9B-KV")
                             steps = gr.Slider(label="Steps", minimum=1, maximum=50, value=4, step=1)
-                            upscale_factor = gr.Dropdown(
-                                label="Upscale model",
-                                info="Applied after generation. UltraSharp / Remacri are sharper than RealESRGAN for clean photo inputs.",
-                                choices=list(UPSCALE_MODELS.keys()),
-                                value="None",
+                            upscale_factor = gr.Dropdown(label="Upscale model",
+                                                         choices=list(UPSCALE_MODELS.keys()), value="None")
+                            gr.Markdown("#### 🖼️ Output canvas size")
+                            canvas_mode = gr.Radio(
+                                choices=["Auto (from base image)", "Custom"],
+                                value="Auto (from base image)", label="Canvas mode",
+                                info=("Auto matches the base image's aspect ratio (longest side 1024). "
+                                      "Use Custom when base and references have very different proportions."),
+                            )
+                            custom_width = gr.Slider(label="Width", minimum=512, maximum=2048,
+                                                     step=16, value=1024, visible=False)
+                            custom_height = gr.Slider(label="Height", minimum=512, maximum=2048,
+                                                      step=16, value=1024, visible=False)
+                            canvas_fit_mode = gr.Radio(
+                                choices=["Stretch", "Pad (color)", "Pad (blur)", "Crop (cover)"],
+                                value="Stretch", label="Canvas fit mode",
+                                info=("How input images are placed into the canvas. "
+                                      "Stretch = current default (can squish). "
+                                      "Pad keeps aspect; Crop fills by trimming edges."),
+                            )
+                            pad_color = gr.ColorPicker(label="Pad colour",
+                                                       value="#000000", visible=False)
+
+                            gr.Markdown("#### 🔁 Batch")
+                            batch_count = gr.Slider(label="Number of runs", minimum=1, maximum=12,
+                                                    step=1, value=1)
+                            batch_vary = gr.Radio(
+                                choices=["Random seed each run",
+                                         "Sequential seed (+1 each)",
+                                         "Sweep first LoRA weight"],
+                                value="Random seed each run", label="Variation strategy",
+                                info=("Sweep linearly varies the weight of whichever LoRA is in "
+                                      "slot 1 (first ticked) across the runs."),
                             )
+                            sweep_min = gr.Slider(label="Sweep min weight", minimum=0.0, maximum=2.0,
+                                                  step=0.05, value=0.4, visible=False)
+                            sweep_max = gr.Slider(label="Sweep max weight", minimum=0.0, maximum=2.0,
+                                                  step=0.05, value=1.4, visible=False)
 
-                    # ── Right column ─────────────────────────────────────────────
+                    # ── Right column ─────────────────────────────────────────
                     with gr.Column(scale=1):
-                        output_image = gr.Image(label="Output", interactive=False, format="png", height=320)
-                        used_seed = gr.Textbox(
-                            label="🌱 Seed used",
-                            info="Re-enter this seed in Advanced Settings to reproduce the result exactly.",
-                            interactive=False,
-                            visible=True,
-                            max_lines=1,
+                        # Gallery so batch runs stream in as they finish; type=
+                        # filepath so PNG metadata round-trips to Send→* buttons
+                        # and to the user's downloads.
+                        output_gallery = gr.Gallery(
+                            label="Output", type="filepath", columns=2, rows=2,
+                            height=420, allow_preview=True, preview=True,
+                            object_fit="contain", show_label=True,
+                        )
+                        used_seed = gr.Textbox(label="🌱 Seed used (last run)",
+                                               interactive=False, max_lines=1)
+                        with gr.Row():
+                            send_out_to_base_btn = gr.Button("↩ Send → Base", size="sm")
+                            send_out_to_ref_btn  = gr.Button("↩ Send → Reference", size="sm")
+                        gr.Markdown(
+                            "*Click a gallery thumbnail before Send→ to pick that specific "
+                            "image; otherwise the latest is used. PNGs contain seed, prompt, "
+                            "LoRAs and settings as `parameters` tEXt chunks (sd-webui / "
+                            "Civitai / ComfyUI readable).*"
                         )
 
-                # ── LoRA selector ─────────────────────────────────────────────────
+                # LoRA selector + weight sliders
                 gr.Markdown("### 🎨 Select LoRA(s)")
                 lora_selector = gr.CheckboxGroup(
                     choices=[s["title"] for s in get_selectable_styles({})],
-                    value=[],
-                    label="Active LoRAs — tick one or more",
+                    value=[], label="Active LoRAs — tick one or more",
                 )
-
-                # ── Weight sliders (pre-rendered, shown/hidden dynamically) ───────
                 gr.Markdown("#### Weights for selected LoRAs")
                 weight_sliders = []
                 with gr.Group():
                     for i in range(MAX_LORA_SLOTS):
                         with gr.Row(elem_classes="lora-weight-row"):
-                            s = gr.Slider(
+                            weight_sliders.append(gr.Slider(
                                 minimum=0.0, maximum=2.0, step=0.05, value=1.0,
-                                label=f"LoRA slot {i+1}",
-                                visible=False,
-                                interactive=True,
-                            )
-                            weight_sliders.append(s)
+                                label=f"LoRA slot {i+1}", visible=False, interactive=True,
+                            ))
 
-                # ── Dynamic HuggingFace loader ─────────────────────────────────────
                 with gr.Accordion("➕ Load Custom LoRA from HuggingFace", open=False):
-                    gr.Markdown("Add any FLUX.2-Klein-compatible LoRA from [HuggingFace Hub](https://huggingface.co/models?other=base_model:adapter:black-forest-labs/FLUX.2-klein-9B). "
-                                "*Added LoRAs are only visible in your own session — other visitors won't see them.*")
+                    gr.Markdown("Add any FLUX.2-Klein-compatible LoRA. *Added LoRAs are only visible in your own session.*")
                     with gr.Row():
                         lora_repo_id = gr.Textbox(label="Repo ID", placeholder="username/repo-name")
                     with gr.Row():
-                        lora_weight_name = gr.Textbox(label="Weight filename (optional)", placeholder="pytorch_lora_weights.safetensors")
+                        lora_weight_name = gr.Textbox(label="Weight filename (optional)",
+                                                      placeholder="pytorch_lora_weights.safetensors")
                         lora_adapter_name = gr.Textbox(label="Adapter name (optional)", placeholder="my-lora")
                     with gr.Row():
                         add_lora_btn = gr.Button("Add LoRA", variant="primary")
                         lora_status = gr.Textbox(label="Status", interactive=False)
 
-                # ── Custom Prompt library ──────────────────────────────────────────
                 with gr.Accordion("📝 Custom Prompts", open=False):
-                    gr.Markdown(
-                        "Save reusable prompt snippets for this session. "
-                        "Tick any saved prompt below and it will be appended to your generation alongside any LoRA prompts. "
-                        "*Prompts are stored in memory — they reset when the Space restarts or the worker recycles.*"
-                    )
-
+                    gr.Markdown("Save reusable prompt snippets for this session.")
                     custom_prompt_selector = gr.CheckboxGroup(
-                        choices=[],
-                        value=[],
+                        choices=[], value=[],
                         label="Saved prompts — tick to append to generation",
                     )
-
                     with gr.Row():
-                        prompt_name_input = gr.Textbox(
-                            label="Name",
-                            placeholder="e.g. Skin detail enhancer",
-                            scale=1,
-                        )
+                        prompt_name_input = gr.Textbox(label="Name",
+                                                       placeholder="e.g. Skin detail enhancer", scale=1)
                     with gr.Row():
-                        prompt_text_input = gr.Textbox(
-                            label="Prompt text",
-                            placeholder="Enter the prompt snippet you want to save and reuse…",
-                            lines=4,
-                        )
+                        prompt_text_input = gr.Textbox(label="Prompt text", lines=4,
+                                                       placeholder="Enter the prompt snippet you want to save…")
                     with gr.Row():
                         add_prompt_btn = gr.Button("💾 Save Prompt", variant="primary")
                         prompt_status = gr.Textbox(label="Status", interactive=False, scale=2)
                     with gr.Row():
-                        delete_prompt_name = gr.Dropdown(
-                            label="Delete a saved prompt",
-                            choices=[],
-                            value=None,
-                            interactive=True,
-                            scale=2,
-                        )
+                        delete_prompt_name = gr.Dropdown(label="Delete a saved prompt",
+                                                         choices=[], value=None, interactive=True, scale=2)
                         delete_prompt_btn = gr.Button("🗑️ Delete", variant="secondary", scale=1)
 
-                gr.Markdown("*Experimental space — results may vary when combining multiple LoRAs.*")
-
-            # ── Crop / Fix Image tab ─────���───────────────────────────────────────
+            # ── Crop / Fix Image tab ─────────────────────────────────────────
             with gr.Tab("✂️ Crop / Fix Image", id="tab_editor"):
                 gr.Markdown(
-                    "Upload an image to crop it, then send the result to the Base Image or add it as a "
-                    "Reference image on the Generate tab. EXIF orientation (the most common cause of "
-                    "images coming out sideways) is corrected automatically as soon as you upload here."
+                    "Upload an image to crop / paint on it, then send the result to the Base "
+                    "Image or add it as a Reference. EXIF orientation is corrected on export."
                 )
                 editor = gr.ImageEditor(
-                    label="Editor",
-                    type="pil",
-                    transforms=["crop"],
-                    brush=False,
-                    eraser=False,
-                    layers=False,
+                    label="Editor", type="pil",
+                    transforms=("crop",),
+                    brush=gr.Brush(default_size=12,
+                                   colors=["#FF4500", "#FFFFFF", "#000000",
+                                           "#FF0000", "#00FF00", "#0000FF"],
+                                   color_mode="fixed"),
+                    eraser=gr.Eraser(default_size=20),
+                    layers=True,
                     sources=["upload", "clipboard"],
                     height=420,
                 )
+                with gr.Row():
+                    heic_uploader = gr.File(
+                        label="📸 Load HEIC / HEIF (iPhone photos)",
+                        file_types=[".heic", ".heif", ".HEIC", ".HEIF"],
+                        file_count="single", type="filepath",
+                    )
                 with gr.Row():
                     send_to_base_btn = gr.Button("→ Send to Base Image", variant="primary")
-                    send_to_ref_btn = gr.Button("→ Add to Reference Images")
+                    send_to_ref_btn  = gr.Button("→ Add to Reference Images")
 
-    # ── Event wiring ──────────────────────────────────────────────────────────
+            # ── Bulk processing tab ──────────────────────────────────────────
+            with gr.Tab("📦 Bulk Process", id="tab_bulk"):
+                gr.Markdown(
+                    "Upload many images and process them with the **same settings as the "
+                    "Generate tab** (prompt, LoRAs, weights, canvas, upscaler, etc.). "
+                    "Outputs stream in one-by-one — each image is its own GPU call, so a "
+                    "ZeroGPU quota wall mid-run only loses the in-progress item. "
+                    "Earlier outputs stay in the gallery and on disk under `/tmp/bulk_<id>/`."
+                )
+                bulk_files = gr.File(
+                    label="Input images",
+                    file_count="multiple", type="filepath",
+                    file_types=["image", ".heic", ".heif"],
+                )
+                with gr.Row():
+                    bulk_run_btn  = gr.Button("▶ Start bulk run", variant="primary")
+                    bulk_stop_btn = gr.Button("⏹ Stop", variant="stop")
+                bulk_status = gr.Markdown("*Ready.*")
+                bulk_gallery = gr.Gallery(
+                    label="Bulk outputs", type="filepath",
+                    columns=4, rows=2, height=480, allow_preview=True, object_fit="contain",
+                )
+                bulk_zip = gr.File(label="📥 Download all (zip + manifest.csv)",
+                                   interactive=False)
 
-    base_image.change(
-        fn=on_base_image_change,
-        inputs=[base_image],
-        outputs=[size_info],
-    )
+            # ── Depth / Pose tab ─────────────────────────────────────────────
+            with gr.Tab("🦴 Depth / Pose", id="tab_control"):
+                gr.Markdown(
+                    "Generate ControlNet-style **depthmaps** and editable **OpenPose** "
+                    "skeletons. The result feeds well into the **RefControl – Depth** / "
+                    "**RefControl – Pose** LoRAs on the Generate tab when sent as a "
+                    "Reference image."
+                )
 
-    reference_images.change(
-        fn=on_reference_change,
-        inputs=[reference_images],
-        outputs=[reference_info],
-    )
+                # Per-tab state — kept here, not in module globals, so each
+                # session edits its own pose without crosstalk.
+                pose_source_state = gr.State(None)        # PIL of last source image
+                pose_keypoints_state = gr.State([])       # list[list[dict]]
+
+                with gr.Row():
+                    with gr.Column(scale=1):
+                        ctrl_source = gr.Image(
+                            label="Source image", type="pil",
+                            sources=["upload", "clipboard"], height=320,
+                        )
+                        with gr.Row():
+                            detect_depth_btn = gr.Button("🌐 Generate depthmap", variant="primary")
+                            detect_pose_btn  = gr.Button("🦴 Detect pose", variant="primary")
+                        insert_blank_btn = gr.Button("➕ Insert blank skeleton template")
+
+                    with gr.Column(scale=1):
+                        depth_output = gr.Image(label="Depthmap", type="pil",
+                                                interactive=False, height=320, format="png")
+                        with gr.Row():
+                            send_depth_ref_btn  = gr.Button("→ Send depth to Reference",
+                                                            variant="primary")
+                            send_depth_base_btn = gr.Button("→ Send depth to Base")
+
+                gr.Markdown("### ✏️ Pose editor")
+                gr.Markdown(
+                    "Pick a person and a joint, then **click anywhere on the editor preview** "
+                    "to move that joint. Hidden joints can be re-added the same way — select "
+                    "them and click. Use the buttons below for delete / clear / re-detect."
+                )
+
+                with gr.Row():
+                    with gr.Column(scale=1):
+                        # interactive=False so users can't accidentally upload a new image
+                        # into this slot. .select still fires for click coordinates.
+                        pose_overlay = gr.Image(
+                            label="Editor — click to place active joint",
+                            type="pil", interactive=False, height=420, format="png",
+                        )
+                    with gr.Column(scale=1):
+                        pose_clean = gr.Image(
+                            label="Skeleton (sent to Reference / Base)",
+                            type="pil", interactive=False, height=420, format="png",
+                        )
+
+                with gr.Row():
+                    active_person_dd = gr.Dropdown(
+                        label="Active person", choices=[], value=None, interactive=True,
+                    )
+                    active_joint_dd = gr.Dropdown(
+                        label="Active joint",
+                        choices=list(OPENPOSE_KEYPOINT_NAMES),
+                        value=None, interactive=True,
+                    )
+
+                with gr.Row():
+                    delete_joint_btn = gr.Button("🗑️ Hide active joint")
+                    reset_pose_btn   = gr.Button("🔄 Re-detect from source")
+                    clear_pose_btn   = gr.Button("🧹 Clear all joints")
+
+                with gr.Row():
+                    send_pose_ref_btn  = gr.Button("→ Send pose to Reference",
+                                                   variant="primary")
+                    send_pose_base_btn = gr.Button("→ Send pose to Base")
+                  
+    # ── Event wiring ─────────────────────────────────────────────────────────
+
+    # HEIC preview fix on the main tab
+    base_image.upload(fn=reencode_upload, inputs=[base_image], outputs=[base_image])
+
+    base_image.change(fn=on_base_image_change, inputs=[base_image], outputs=[size_info])
+    reference_images.change(fn=on_reference_change, inputs=[reference_images], outputs=[reference_info])
 
     lora_selector.change(
         fn=update_weight_sliders,
@@ -1025,53 +990,259 @@ with gr.Blocks(css=css, theme=orange_red_theme) as demo:
     add_prompt_btn.click(
         fn=add_custom_prompt,
         inputs=[prompt_name_input, prompt_text_input, custom_prompts_state, custom_prompt_counter_state],
-        outputs=[prompt_status, custom_prompts_state, custom_prompt_counter_state, custom_prompt_selector, delete_prompt_name, prompt_name_input],
+        outputs=[prompt_status, custom_prompts_state, custom_prompt_counter_state,
+                 custom_prompt_selector, delete_prompt_name, prompt_name_input],
     )
-
     delete_prompt_btn.click(
         fn=delete_custom_prompt,
         inputs=[delete_prompt_name, custom_prompt_selector, custom_prompts_state],
         outputs=[prompt_status, custom_prompts_state, custom_prompt_selector, delete_prompt_name],
     )
-
     custom_prompt_selector.change(
         fn=update_custom_prompt_display,
         inputs=[custom_prompt_selector, custom_prompts_state],
         outputs=[custom_prompt_display],
     )
 
-    run_button.click(
+    # Canvas / fit / batch UI toggles
+    canvas_mode.change(fn=on_canvas_mode_change, inputs=[canvas_mode],
+                       outputs=[custom_width, custom_height])
+    canvas_fit_mode.change(fn=on_fit_mode_change, inputs=[canvas_fit_mode],
+                           outputs=[pad_color])
+    batch_vary.change(fn=on_batch_vary_change, inputs=[batch_vary],
+                      outputs=[sweep_min, sweep_max])
+
+    # Track which gallery item the user clicked, so Send→Base/Ref can use it
+    output_gallery.select(fn=on_gallery_select, inputs=[output_gallery],
+                          outputs=[selected_output_state])
+
+    # The Generate-tab generator. .click() returns an event we keep so the
+    # bulk Stop button can cancel mid-stream too if desired.
+    run_event = run_button.click(
         fn=infer,
-        inputs=[base_image, reference_images, prompt, lora_prompt_display, custom_prompt_display, lora_selector, seed, randomize_seed, guidance_scale, steps, upscale_factor, dynamic_loras_state] + weight_sliders,
-        outputs=[output_image, used_seed],
+        inputs=[base_image, reference_images, prompt, lora_prompt_display, custom_prompt_display,
+                lora_selector, seed, randomize_seed, guidance_scale, steps, upscale_factor,
+                canvas_mode, custom_width, custom_height, canvas_fit_mode, pad_color,
+                batch_count, batch_vary, sweep_min, sweep_max,
+                dynamic_loras_state] + weight_sliders,
+        outputs=[output_gallery, used_seed],
     )
 
-    # ── Image Editor tab wiring ─────────────────────────────────────────────
-    editor.upload(
-        fn=fix_editor_orientation,
-        inputs=[editor],
-        outputs=[editor],
-    )
+    # ── Editor tab wiring ────────────────────────────────────────────────────
+    # NOTE: do NOT add editor.upload(outputs=[editor]) — remounts the editor.
+
+    heic_uploader.upload(fn=load_heic_to_editor, inputs=[heic_uploader], outputs=[editor])
 
-    send_to_base_btn.click(
-        fn=send_editor_to_base,
-        inputs=[editor],
+    send_to_base_btn.click(fn=send_editor_to_base, inputs=[editor], outputs=[base_image]) \
+        .then(fn=on_base_image_change, inputs=[base_image], outputs=[size_info]) \
+        .then(fn=lambda: gr.Tabs(selected="tab_generate"), outputs=[main_tabs])
+
+    send_to_ref_btn.click(fn=send_editor_to_reference,
+                          inputs=[editor, reference_images], outputs=[reference_images]) \
+        .then(fn=on_reference_change, inputs=[reference_images], outputs=[reference_info]) \
+        .then(fn=lambda: gr.Tabs(selected="tab_generate"), outputs=[main_tabs])
+
+    # Send-output buttons use the selected gallery item (or latest fallback)
+    send_out_to_base_btn.click(
+        fn=send_output_to_base,
+        inputs=[selected_output_state, output_gallery],
         outputs=[base_image],
-    ).then(
-        fn=on_base_image_change, inputs=[base_image], outputs=[size_info],
-    ).then(
-        fn=lambda: gr.Tabs(selected="tab_generate"), outputs=[main_tabs],
-    )
+    ).then(fn=on_base_image_change, inputs=[base_image], outputs=[size_info])
 
-    send_to_ref_btn.click(
-        fn=send_editor_to_reference,
-        inputs=[editor, reference_images],
+    send_out_to_ref_btn.click(
+        fn=send_output_to_reference,
+        inputs=[selected_output_state, output_gallery, reference_images],
         outputs=[reference_images],
-    ).then(
-        fn=on_reference_change, inputs=[reference_images], outputs=[reference_info],
-    ).then(
-        fn=lambda: gr.Tabs(selected="tab_generate"), outputs=[main_tabs],
+    ).then(fn=on_reference_change, inputs=[reference_images], outputs=[reference_info])
+
+    # ── Bulk tab wiring ──────────────────────────────────────────────────────
+    # Reuses Generate-tab settings as inputs verbatim — single source of truth,
+    # no two-way state sync to keep in step.
+    bulk_event = bulk_run_btn.click(
+        fn=bulk_infer,
+        inputs=[bulk_files,
+                prompt, lora_prompt_display, custom_prompt_display, lora_selector,
+                seed, randomize_seed, guidance_scale, steps, upscale_factor,
+                canvas_mode, custom_width, custom_height, canvas_fit_mode, pad_color,
+                dynamic_loras_state] + weight_sliders,
+        outputs=[bulk_gallery, bulk_status, bulk_zip],
+    )
+
+    # Stop cancels both running generators; the in-flight GPU call still
+    # finishes (ZeroGPU can't be killed mid-step), but no further iterations
+    # start. Any already-saved outputs remain in the gallery and on disk.
+    bulk_stop_btn.click(fn=lambda: gr.Info("Stop requested — finishing current image."),
+                        cancels=[bulk_event, run_event])
+
+    # ── Depth / Pose tab wiring ──────────────────────────────────────────────
+
+    # Cache the last successfully-loaded source so re-renders after edits
+    # don't need the user to keep the upload widget populated.
+    ctrl_source.change(
+        fn=lambda img: img, inputs=[ctrl_source], outputs=[pose_source_state],
+    )
+
+    # Depth generation
+    detect_depth_btn.click(
+        fn=generate_depthmap, inputs=[ctrl_source], outputs=[depth_output],
     )
 
+    # Pose detection → fills state, dropdowns, both preview images.
+    def _on_detect_pose(source):
+        if source is None:
+            raise gr.Error("Upload a source image first.")
+        poses, w, h = detect_pose(source)
+        if not poses:
+            gr.Warning("No people detected — try 'Insert blank skeleton template' "
+                       "or a different image.")
+            return ([], gr.update(choices=[], value=None),
+                    gr.update(value=None), None, None)
+        return (
+            poses,
+            gr.update(choices=person_choices(poses), value="Person 1"),
+            gr.update(value=OPENPOSE_KEYPOINT_NAMES[0]),
+            render_pose_overlay(source, poses, 0, 0),
+            render_pose_skeleton(poses, w, h),
+        )
+
+    detect_pose_btn.click(
+        fn=_on_detect_pose, inputs=[ctrl_source],
+        outputs=[pose_keypoints_state, active_person_dd, active_joint_dd,
+                 pose_overlay, pose_clean],
+    )
+    reset_pose_btn.click(  # same handler — re-runs detection
+        fn=_on_detect_pose, inputs=[ctrl_source],
+        outputs=[pose_keypoints_state, active_person_dd, active_joint_dd,
+                 pose_overlay, pose_clean],
+    )
+
+    # Insert a default standing-figure template centred in the source canvas.
+    def _on_insert_blank(source):
+        if source is None:
+            raise gr.Error("Upload a source image first.")
+        w, h = source.size
+        poses = [default_pose_template(w, h)]
+        return (
+            poses,
+            gr.update(choices=["Person 1"], value="Person 1"),
+            gr.update(value=OPENPOSE_KEYPOINT_NAMES[0]),
+            render_pose_overlay(source, poses, 0, 0),
+            render_pose_skeleton(poses, w, h),
+        )
+
+    insert_blank_btn.click(
+        fn=_on_insert_blank, inputs=[ctrl_source],
+        outputs=[pose_keypoints_state, active_person_dd, active_joint_dd,
+                 pose_overlay, pose_clean],
+    )
+
+    # ── Click-to-edit: the heart of the pose editor ──
+    # gr.SelectData on a gr.Image gives .index = (x, y) in image pixels, even
+    # when interactive=False, which is exactly what we need.
+    def _on_overlay_click(evt: gr.SelectData, poses, source, person_label, joint_name):
+        if not poses or source is None or evt is None or evt.index is None:
+            return gr.update(), gr.update(), gr.update()
+        person_idx = parse_person_idx(person_label)
+        joint_idx  = joint_name_to_index(joint_name)
+        if person_idx is None or joint_idx < 0:
+            return gr.update(), gr.update(), gr.update()
+        x, y = evt.index
+        w, h = source.size
+        new_poses = move_joint(poses, person_idx, joint_idx, x, y, w, h)
+        return (
+            new_poses,
+            render_pose_overlay(source, new_poses, person_idx, joint_idx),
+            render_pose_skeleton(new_poses, w, h),
+        )
+
+    pose_overlay.select(
+        fn=_on_overlay_click,
+        inputs=[pose_keypoints_state, pose_source_state,
+                active_person_dd, active_joint_dd],
+        outputs=[pose_keypoints_state, pose_overlay, pose_clean],
+    )
+
+    # Changing the active joint or person just re-renders the overlay so the
+    # highlight ring follows — keypoints are not mutated.
+    def _on_active_change(poses, source, person_label, joint_name):
+        if not poses or source is None:
+            return gr.update()
+        person_idx = parse_person_idx(person_label) or 0
+        joint_idx  = max(joint_name_to_index(joint_name), 0)
+        return render_pose_overlay(source, poses, person_idx, joint_idx)
+
+    active_person_dd.change(
+        fn=_on_active_change,
+        inputs=[pose_keypoints_state, pose_source_state,
+                active_person_dd, active_joint_dd],
+        outputs=[pose_overlay],
+    )
+    active_joint_dd.change(
+        fn=_on_active_change,
+        inputs=[pose_keypoints_state, pose_source_state,
+                active_person_dd, active_joint_dd],
+        outputs=[pose_overlay],
+    )
+
+    # Hide / clear
+    def _on_hide_active(poses, source, person_label, joint_name):
+        person_idx = parse_person_idx(person_label)
+        joint_idx  = joint_name_to_index(joint_name)
+        new_poses  = hide_joint(poses, person_idx, joint_idx)
+        if source is None:
+            return new_poses, gr.update(), gr.update()
+        w, h = source.size
+        return (new_poses,
+                render_pose_overlay(source, new_poses, person_idx, joint_idx),
+                render_pose_skeleton(new_poses, w, h))
+
+    delete_joint_btn.click(
+        fn=_on_hide_active,
+        inputs=[pose_keypoints_state, pose_source_state,
+                active_person_dd, active_joint_dd],
+        outputs=[pose_keypoints_state, pose_overlay, pose_clean],
+    )
+
+    def _on_clear_all(poses, source):
+        new_poses = clear_all_joints(poses)
+        if source is None:
+            return new_poses, gr.update(), gr.update()
+        w, h = source.size
+        return (new_poses,
+                render_pose_overlay(source, new_poses, None, None),
+                render_pose_skeleton(new_poses, w, h))
+
+    clear_pose_btn.click(
+        fn=_on_clear_all,
+        inputs=[pose_keypoints_state, pose_source_state],
+        outputs=[pose_keypoints_state, pose_overlay, pose_clean],
+    )
+
+    # Send → main tab. We reuse the existing base_image / reference_images
+    # components so users land back on the Generate tab fully wired up.
+    send_depth_ref_btn.click(
+        fn=push_pil_to_reference, inputs=[depth_output, reference_images],
+        outputs=[reference_images],
+    ).then(fn=on_reference_change, inputs=[reference_images], outputs=[reference_info]
+    ).then(fn=lambda: gr.Tabs(selected="tab_generate"), outputs=[main_tabs])
+
+    send_depth_base_btn.click(
+        fn=push_pil_to_base, inputs=[depth_output], outputs=[base_image],
+    ).then(fn=on_base_image_change, inputs=[base_image], outputs=[size_info]
+    ).then(fn=lambda: gr.Tabs(selected="tab_generate"), outputs=[main_tabs])
+
+    send_pose_ref_btn.click(
+        fn=push_pil_to_reference, inputs=[pose_clean, reference_images],
+        outputs=[reference_images],
+    ).then(fn=on_reference_change, inputs=[reference_images], outputs=[reference_info]
+    ).then(fn=lambda: gr.Tabs(selected="tab_generate"), outputs=[main_tabs])
+
+    send_pose_base_btn.click(
+        fn=push_pil_to_base, inputs=[pose_clean], outputs=[base_image],
+    ).then(fn=on_base_image_change, inputs=[base_image], outputs=[size_info]
+    ).then(fn=lambda: gr.Tabs(selected="tab_generate"), outputs=[main_tabs])
+
 if __name__ == "__main__":
-    demo.queue().launch(css=css, theme=orange_red_theme, mcp_server=True, ssr_mode=False, show_error=True)
+    # Gradio 6.0: theme and css go here, not on Blocks()
+    demo.queue().launch(css=css, theme=orange_red_theme,
+                        mcp_server=True, ssr_mode=False, show_error=True)

config.py

@@ -0,0 +1,28 @@
+"""Environment-derived configuration and global constants.
+
+Importing this module has no side effects beyond reading env vars and computing
+a handful of constants. It must not import torch, diffusers, or gradio so it can
+be imported anywhere without dragging in heavy dependencies.
+"""
+import os
+import numpy as np
+
+# ── Model variant selection ──────────────────────────────────────────────────
+MODEL_VARIANT = os.environ.get("MODEL_VARIANT", "9B")
+if MODEL_VARIANT == "9B-KV":
+    _MODEL_REPO = "black-forest-labs/FLUX.2-klein-9b-kv"
+else:
+    MODEL_VARIANT = "9B"  # normalise any typos back to default
+    _MODEL_REPO = "black-forest-labs/FLUX.2-klein-9B"
+
+MODEL_REPO = _MODEL_REPO
+
+# ── Misc constants ───────────────────────────────────────────────────────────
+MAX_SEED = int(np.iinfo(np.int32).max)
+
+# Maximum number of simultaneous LoRA weight sliders to pre-render in the UI.
+MAX_LORA_SLOTS = 6
+
+# Logging toggle — read once, exposed as a plain bool so other modules don't
+# repeat the env-string-parsing dance.
+ENABLE_LOGGING = os.environ.get("ENABLE_LOGGING", "No").strip().lower() == "yes"

control_tools.py

@@ -0,0 +1,269 @@
+"""Depth + pose estimation helpers for the Depth/Pose tab.
+
+Two heavy models are lazy-loaded on first use, both inside @spaces.GPU
+functions so ZeroGPU schedules them as normal short-lived workloads:
+  - Depth-Anything-V2-Small via transformers' depth-estimation pipeline
+  - OpenposeDetector (lllyasviel/Annotators) via controlnet_aux
+
+Pose keypoints round-trip through a plain JSON-friendly list-of-dicts so they
+serialise cleanly into gr.State — that's what makes click-to-edit feasible.
+"""
+from __future__ import annotations
+
+import numpy as np
+from PIL import Image, ImageDraw, ImageEnhance
+import gradio as gr
+import torch
+import spaces
+
+
+# ── Standard OpenPose body-18 layout (COCO-18 ordering) ─────────────────────
+OPENPOSE_KEYPOINT_NAMES = [
+    "nose", "neck",
+    "right_shoulder", "right_elbow", "right_wrist",
+    "left_shoulder", "left_elbow", "left_wrist",
+    "right_hip", "right_knee", "right_ankle",
+    "left_hip", "left_knee", "left_ankle",
+    "right_eye", "left_eye", "right_ear", "left_ear",
+]
+
+# Bones as 0-indexed (a, b) pairs into OPENPOSE_KEYPOINT_NAMES. Mirrors the
+# OpenPose `limbSeq` constant so ControlNet/RefControl-Pose readers see the
+# exact skeleton topology they expect.
+SKELETON_EDGES = [
+    (1, 2), (1, 5), (2, 3), (3, 4), (5, 6), (6, 7),
+    (1, 8), (8, 9), (9, 10), (1, 11), (11, 12), (12, 13),
+    (1, 0), (0, 14), (14, 16), (0, 15), (15, 17),
+]
+
+# OpenPose-canonical 18-joint colour table (used for both bones and dots).
+JOINT_COLORS = [
+    (255, 0, 0), (255, 85, 0), (255, 170, 0), (255, 255, 0), (170, 255, 0),
+    (85, 255, 0), (0, 255, 0), (0, 255, 85), (0, 255, 170), (0, 255, 255),
+    (0, 170, 255), (0, 85, 255), (0, 0, 255), (85, 0, 255), (170, 0, 255),
+    (255, 0, 255), (255, 0, 170), (255, 0, 85),
+]
+
+
+# ── Lazy model loaders ───────────────────────────────────────────────────────
+
+_depth_pipeline = None
+_pose_detector = None
+
+
+def _load_depth_pipeline():
+    """Loaded inside the first @spaces.GPU call, then cached for the rest of
+    the worker's lifetime. Avoids paying the cold-start cost for users who
+    only ever use depth or only ever use pose."""
+    global _depth_pipeline
+    if _depth_pipeline is None:
+        from transformers import pipeline
+        _depth_pipeline = pipeline(
+            "depth-estimation",
+            model="depth-anything/Depth-Anything-V2-Small-hf",
+            device=0 if torch.cuda.is_available() else -1,
+        )
+    return _depth_pipeline
+
+
+def _load_pose_detector():
+    global _pose_detector
+    if _pose_detector is None:
+        try:
+            from controlnet_aux import OpenposeDetector
+        except ImportError:
+            raise gr.Error(
+                "controlnet_aux is not installed. Add `controlnet_aux` to requirements.txt."
+            )
+        _pose_detector = OpenposeDetector.from_pretrained("lllyasviel/Annotators")
+        if torch.cuda.is_available():
+            _pose_detector = _pose_detector.to("cuda")
+    return _pose_detector
+
+
+# ── Depth ────────────────────────────────────────────────────────────────────
+
+@spaces.GPU
+def generate_depthmap(image: Image.Image) -> Image.Image:
+    if image is None:
+        raise gr.Error("Upload a source image first.")
+    pipe = _load_depth_pipeline()
+    depth = pipe(image.convert("RGB"))["depth"]  # PIL L
+    return depth.convert("RGB")
+
+
+# ── Pose detection → editable keypoint list ─────────────────────────────────
+
+@spaces.GPU
+def detect_pose(image: Image.Image) -> tuple[list[list[dict]], int, int]:
+    """Run pose detection and return:
+      - poses : list of dicts, one per detected person. Each pose is a list of
+                18 entries {"name", "x", "y", "visible"} — in *pixel* coords
+                (not normalised), so the editor can pass click positions in
+                directly without a coordinate transform.
+      - w, h  : source image dimensions.
+    """
+    if image is None:
+        raise gr.Error("Upload a source image first.")
+    img = image.convert("RGB")
+    w, h = img.size
+    detector = _load_pose_detector()
+    try:
+        poses_raw = detector.detect_poses(np.array(img))
+    except Exception as e:
+        raise gr.Error(f"Pose detection failed: {e}")
+
+    out = []
+    for pose in poses_raw:
+        body_kps = pose.body.keypoints if pose.body else [None] * 18
+        person = []
+        for i in range(18):
+            kp = body_kps[i] if i < len(body_kps) else None
+            name = OPENPOSE_KEYPOINT_NAMES[i]
+            # controlnet_aux returns normalised (x, y) ∈ [0, 1] with a score.
+            # Drop low-confidence detections so the editor starts clean.
+            if kp is None or (getattr(kp, "score", 1.0) or 0) < 0.3:
+                person.append({"name": name, "x": 0.0, "y": 0.0, "visible": False})
+            else:
+                person.append({
+                    "name": name,
+                    "x": float(kp.x) * w,
+                    "y": float(kp.y) * h,
+                    "visible": True,
+                })
+        out.append(person)
+    return out, w, h
+
+
+# ── Rendering ───────────────────────────────────────────────────────────────
+
+def render_pose_skeleton(poses: list[list[dict]], width: int, height: int) -> Image.Image:
+    """Skeleton on a black canvas — this is what ControlNet/RefControl-Pose wants."""
+    canvas = Image.new("RGB", (width, height), (0, 0, 0))
+    draw = ImageDraw.Draw(canvas)
+    for pose in poses:
+        for a, b in SKELETON_EDGES:
+            ka, kb = pose[a], pose[b]
+            if not (ka["visible"] and kb["visible"]):
+                continue
+            draw.line([ka["x"], ka["y"], kb["x"], kb["y"]],
+                      fill=JOINT_COLORS[a], width=4)
+        for i, kp in enumerate(pose):
+            if not kp["visible"]:
+                continue
+            r = 4
+            draw.ellipse([kp["x"] - r, kp["y"] - r, kp["x"] + r, kp["y"] + r],
+                         fill=JOINT_COLORS[i])
+    return canvas
+
+
+def render_pose_overlay(
+    source: Image.Image,
+    poses: list[list[dict]],
+    active_person: int | None = None,
+    active_joint: int | None = None,
+) -> Image.Image | None:
+    """Skeleton on top of a dimmed source — the editing view. The active joint
+    gets a white-ringed highlight so the user can see what their next click
+    will move."""
+    if source is None:
+        return None
+    w, h = source.size
+    base = ImageEnhance.Brightness(source.convert("RGB")).enhance(0.45)
+    skel = render_pose_skeleton(poses, w, h)
+    np_base, np_skel = np.array(base), np.array(skel)
+    mask = np_skel.any(axis=2)
+    np_base[mask] = np_skel[mask]
+    out = Image.fromarray(np_base)
+
+    if (active_person is not None and 0 <= active_person < len(poses)
+            and active_joint is not None and 0 <= active_joint < 18):
+        kp = poses[active_person][active_joint]
+        if kp["visible"]:
+            d = ImageDraw.Draw(out)
+            x, y = kp["x"], kp["y"]
+            # Double ring so it stays visible on any background colour
+            for r, col in [(10, (255, 255, 255)), (7, (0, 0, 0))]:
+                d.ellipse([x - r, y - r, x + r, y + r], outline=col, width=2)
+    return out
+
+
+# ── Edit operations ─────────────────────────────────────────────────────────
+
+def move_joint(poses, person_idx, joint_idx, x, y, width, height):
+    """Move (and make visible) a joint. Clamped to image bounds so clicks just
+    outside the canvas don't fly off."""
+    if (not poses or person_idx is None or joint_idx is None
+            or not (0 <= person_idx < len(poses) and 0 <= joint_idx < 18)):
+        return poses
+    x = float(max(0, min(x, width - 1)))
+    y = float(max(0, min(y, height - 1)))
+    new = [list(p) for p in poses]
+    new[person_idx][joint_idx] = {
+        **new[person_idx][joint_idx], "x": x, "y": y, "visible": True,
+    }
+    return new
+
+
+def hide_joint(poses, person_idx, joint_idx):
+    if (not poses or person_idx is None or joint_idx is None
+            or not (0 <= person_idx < len(poses) and 0 <= joint_idx < 18)):
+        return poses
+    new = [list(p) for p in poses]
+    new[person_idx][joint_idx] = {**new[person_idx][joint_idx], "visible": False}
+    return new
+
+
+def clear_all_joints(poses):
+    return [[{**kp, "visible": False} for kp in pose] for pose in (poses or [])]
+
+
+def default_pose_template(width: int, height: int) -> list[dict]:
+    """Standing-figure skeleton centred in the canvas — handy when detection
+    misses everyone, or when you want to draw a pose from scratch."""
+    cx = width / 2
+    s = min(width, height) * 0.40
+    top = height / 2 - s
+    def y(f): return top + f * (2 * s)
+    return [
+        {"name": "nose",           "x": cx,             "y": y(0.05), "visible": True},
+        {"name": "neck",           "x": cx,             "y": y(0.15), "visible": True},
+        {"name": "right_shoulder", "x": cx + s * 0.18,  "y": y(0.18), "visible": True},
+        {"name": "right_elbow",    "x": cx + s * 0.25,  "y": y(0.35), "visible": True},
+        {"name": "right_wrist",    "x": cx + s * 0.30,  "y": y(0.50), "visible": True},
+        {"name": "left_shoulder",  "x": cx - s * 0.18,  "y": y(0.18), "visible": True},
+        {"name": "left_elbow",     "x": cx - s * 0.25,  "y": y(0.35), "visible": True},
+        {"name": "left_wrist",     "x": cx - s * 0.30,  "y": y(0.50), "visible": True},
+        {"name": "right_hip",      "x": cx + s * 0.12,  "y": y(0.55), "visible": True},
+        {"name": "right_knee",     "x": cx + s * 0.13,  "y": y(0.75), "visible": True},
+        {"name": "right_ankle",    "x": cx + s * 0.14,  "y": y(0.95), "visible": True},
+        {"name": "left_hip",       "x": cx - s * 0.12,  "y": y(0.55), "visible": True},
+        {"name": "left_knee",      "x": cx - s * 0.13,  "y": y(0.75), "visible": True},
+        {"name": "left_ankle",     "x": cx - s * 0.14,  "y": y(0.95), "visible": True},
+        {"name": "right_eye",      "x": cx + s * 0.025, "y": y(0.03), "visible": True},
+        {"name": "left_eye",       "x": cx - s * 0.025, "y": y(0.03), "visible": True},
+        {"name": "right_ear",      "x": cx + s * 0.05,  "y": y(0.05), "visible": True},
+        {"name": "left_ear",       "x": cx - s * 0.05,  "y": y(0.05), "visible": True},
+    ]
+
+
+# ── Dropdown helpers ─────────────────────────────────────────────────────────
+
+def person_choices(poses):
+    return [f"Person {i+1}" for i in range(len(poses or []))]
+
+def parse_person_idx(label: str | None) -> int | None:
+    if not label:
+        return None
+    try:
+        return int(label.replace("Person ", "")) - 1
+    except ValueError:
+        return None
+
+def joint_name_to_index(name: str | None) -> int:
+    if not name:
+        return -1
+    try:
+        return OPENPOSE_KEYPOINT_NAMES.index(name)
+    except ValueError:
+        return -1

image_utils.py

@@ -0,0 +1,318 @@
+"""Pure image helpers — no torch, no diffusers, no gradio state.
+
+Owns: EXIF handling, dimension snapping, canvas fitting, editor-composite
+extraction, HEIC decoding, PNG metadata embedding.
+"""
+from __future__ import annotations
+
+import json
+import tempfile
+from typing import Any
+
+import numpy as np
+from PIL import Image, ImageOps, ImageFilter
+from PIL.PngImagePlugin import PngInfo
+import gradio as gr
+
+# ── HEIC / HEIF support ──────────────────────────────────────────────────────
+try:
+    from pillow_heif import register_heif_opener
+    register_heif_opener()
+except ImportError:
+    print("pillow-heif not installed — HEIC/HEIF uploads will not work. "
+          "Add `pillow-heif` to requirements.txt.")
+
+
+# ── EXIF / dimension helpers ─────────────────────────────────────────────────
+
+def fix_orientation(img: Image.Image | None) -> Image.Image | None:
+    if img is None:
+        return None
+    return ImageOps.exif_transpose(img)
+
+
+def _snap16(v: float) -> int:
+    """Snap to a multiple of 16 — required by FLUX's VAE."""
+    return max(16, (int(v) // 16) * 16)
+
+
+def compute_base_dimensions(image: Image.Image | None) -> tuple[int, int]:
+    if image is None:
+        return 1024, 1024
+    w, h = image.size
+    scale = min(1024 / w, 1024 / h)
+    return _snap16(w * scale), _snap16(h * scale)
+
+
+update_dimensions_on_upload = compute_base_dimensions
+
+
+def compute_canvas_dimensions(
+    base_image: Image.Image | None,
+    canvas_mode: str,
+    custom_width: int,
+    custom_height: int,
+) -> tuple[int, int]:
+    if canvas_mode == "Custom":
+        return _snap16(custom_width), _snap16(custom_height)
+    return compute_base_dimensions(base_image)
+
+
+# ── Canvas fitting ──────────────────────────────────────────────────────────
+
+def fit_to_canvas(
+    img: Image.Image,
+    width: int,
+    height: int,
+    mode: str = "Stretch",
+    pad_color: str = "#000000",
+) -> Image.Image:
+    """Return `img` resized to exactly width×height using the given strategy.
+
+    Modes:
+      - "Stretch"        : resize ignoring aspect (current default, may distort)
+      - "Pad (color)"    : scale to fit, pad with `pad_color`
+      - "Pad (blur)"     : scale to fit, pad with a blurred cover of the image
+      - "Crop (cover)"   : scale to cover, center-crop to canvas
+    """
+    img = img.convert("RGB")
+
+    if mode == "Stretch":
+        return img.resize((width, height), Image.LANCZOS)
+
+    iw, ih = img.size
+
+    if mode == "Pad (color)":
+        scale = min(width / iw, height / ih)
+        nw, nh = max(1, int(iw * scale)), max(1, int(ih * scale))
+        resized = img.resize((nw, nh), Image.LANCZOS)
+        canvas = Image.new("RGB", (width, height), pad_color)
+        canvas.paste(resized, ((width - nw) // 2, (height - nh) // 2))
+        return canvas
+
+    if mode == "Pad (blur)":
+        # Foreground: scale-to-fit
+        scale = min(width / iw, height / ih)
+        nw, nh = max(1, int(iw * scale)), max(1, int(ih * scale))
+        fg = img.resize((nw, nh), Image.LANCZOS)
+        # Background: scale-to-cover, center-crop, then blur heavily
+        cscale = max(width / iw, height / ih)
+        cw, ch = max(1, int(iw * cscale)), max(1, int(ih * cscale))
+        bg = img.resize((cw, ch), Image.LANCZOS)
+        bg = bg.crop(((cw - width) // 2, (ch - height) // 2,
+                     (cw - width) // 2 + width, (ch - height) // 2 + height))
+        bg = bg.filter(ImageFilter.GaussianBlur(radius=32))
+        bg.paste(fg, ((width - nw) // 2, (height - nh) // 2))
+        return bg
+
+    if mode == "Crop (cover)":
+        cscale = max(width / iw, height / ih)
+        nw, nh = max(1, int(iw * cscale)), max(1, int(ih * cscale))
+        resized = img.resize((nw, nh), Image.LANCZOS)
+        left = (nw - width) // 2
+        top = (nh - height) // 2
+        return resized.crop((left, top, left + width, top + height))
+
+    # Unknown mode → fall back to stretch rather than erroring during inference
+    print(f"[fit_to_canvas] unknown mode {mode!r} — falling back to Stretch.")
+    return img.resize((width, height), Image.LANCZOS)
+
+
+# ── UI label updates ────────────────────────────────────────────────────────
+
+def on_base_image_change(img) -> str:
+    if img is None:
+        return "*No base image uploaded yet*"
+    try:
+        pil_img = img if isinstance(img, Image.Image) else Image.open(img)
+        ow, oh = pil_img.size
+        bw, bh = compute_base_dimensions(pil_img)
+        return (
+            f"Input: **{ow} × {oh}** px → "
+            f"Auto canvas (pre-upscale): **{bw} × {bh}** px"
+        )
+    except Exception as e:
+        return f"*Could not read dimensions: {e}*"
+
+
+def on_reference_change(images) -> str:
+    if not images:
+        return "📷 No reference images"
+    count = len(images)
+    return f"📷 {count} reference image{'s' if count != 1 else ''} uploaded"
+
+
+# ── Upload round-trip (fixes HEIC preview in main tab) ──────────────────────
+
+def reencode_upload(img):
+    if img is None:
+        return None
+    if not isinstance(img, Image.Image):
+        try:
+            img = Image.open(img)
+        except Exception:
+            return img
+    return fix_orientation(img).convert("RGB")
+
+
+# ── Inference input assembly ────────────────────────────────────────────────
+
+def process_images(base_image, reference_images) -> list[Image.Image]:
+    pil_images: list[Image.Image] = []
+    if base_image is not None:
+        try:
+            img = base_image if isinstance(base_image, Image.Image) else Image.open(base_image)
+            pil_images.append(fix_orientation(img).convert("RGB"))
+        except Exception as e:
+            print(f"Skipping invalid base image: {e}")
+    for item in (reference_images or []):
+        try:
+            path_or_img = item[0] if isinstance(item, (tuple, list)) else item
+            if isinstance(path_or_img, Image.Image):
+                img = path_or_img
+            elif isinstance(path_or_img, str):
+                img = Image.open(path_or_img)
+            else:
+                img = Image.open(path_or_img.name)
+            pil_images.append(fix_orientation(img).convert("RGB"))
+        except Exception as e:
+            print(f"Skipping invalid reference image: {e}")
+    return pil_images
+
+
+# ── ImageEditor helpers ─────────────────────────────────────────────────────
+
+def _editor_composite(editor_value) -> Image.Image:
+    if not editor_value or editor_value.get("composite") is None:
+        raise gr.Error("Upload and crop an image in the editor first.")
+    composite = editor_value["composite"]
+    if isinstance(composite, np.ndarray):
+        composite = Image.fromarray(composite)
+    return composite.convert("RGB")
+
+
+def send_editor_to_base(editor_value) -> Image.Image:
+    composite = fix_orientation(_editor_composite(editor_value))
+    gr.Info("Sent to Base Image")
+    return composite
+
+
+def send_editor_to_reference(editor_value, current_gallery) -> list:
+    composite = fix_orientation(_editor_composite(editor_value))
+    current = list(current_gallery or [])
+    current.append(composite)
+    gr.Info("Added to Reference Images")
+    return current
+
+
+def load_heic_to_editor(path):
+    if not path:
+        return gr.update()
+    try:
+        img = fix_orientation(Image.open(path)).convert("RGB")
+    except Exception as e:
+        raise gr.Error(f"Could not decode HEIC/HEIF: {e}")
+    gr.Info("HEIC loaded into editor.")
+    return img
+
+
+# ── Send output → base / reference (gallery-aware) ──────────────────────────
+
+def _resolve_gallery_path(selected_path, gallery_value):
+    """Pick the path the Send-to-* buttons should use.
+
+    Prefers the user's currently-selected gallery item; falls back to the most
+    recent (last) item so single-image runs and "didn't click anything" cases
+    both behave intuitively.
+    """
+    if selected_path:
+        return selected_path
+    if not gallery_value:
+        return None
+    item = gallery_value[-1]
+    return item[0] if isinstance(item, (list, tuple)) else item
+
+
+def send_output_to_base(selected_path, gallery_value):
+    path = _resolve_gallery_path(selected_path, gallery_value)
+    if not path:
+        raise gr.Error("Nothing to send — generate an image first.")
+    img = Image.open(path).convert("RGB")
+    gr.Info("Output sent to Base Image.")
+    return img
+
+
+def send_output_to_reference(selected_path, gallery_value, current_gallery):
+    path = _resolve_gallery_path(selected_path, gallery_value)
+    if not path:
+        raise gr.Error("Nothing to send — generate an image first.")
+    img = Image.open(path).convert("RGB")
+    current = list(current_gallery or [])
+    current.append(img)
+    gr.Info("Output added to Reference Images.")
+    return current
+
+
+# ── PNG metadata embedding ──────────────────────────────────────────────────
+
+def _format_parameters_string(meta: dict[str, Any]) -> str:
+    prompt = meta.get("prompt", "") or ""
+    fields = [
+        ("Seed", meta.get("seed")),
+        ("Steps", meta.get("steps")),
+        ("CFG scale", meta.get("guidance_scale")),
+        ("Size", f"{meta.get('width')}x{meta.get('height')}"),
+        ("Model", meta.get("model")),
+        ("Upscaler", meta.get("upscale_factor")),
+        ("Canvas mode", meta.get("canvas_mode")),
+        ("Fit mode", meta.get("canvas_fit_mode")),
+    ]
+    loras = meta.get("loras") or []
+    if loras:
+        fields.append(("LoRAs", ", ".join(f"{n}:{w:.2f}" for n, w in loras)))
+    kv = ", ".join(f"{k}: {v}" for k, v in fields if v not in (None, "", "None"))
+    return f"{prompt}\n{kv}".strip()
+
+
+def build_pnginfo(meta: dict[str, Any]) -> PngInfo:
+    """Public so bulk processing can reuse it for in-place saves."""
+    info = PngInfo()
+    info.add_text("parameters", _format_parameters_string(meta))
+    for k in ("prompt", "seed", "steps", "guidance_scale", "width", "height",
+              "model", "upscale_factor", "canvas_mode", "canvas_fit_mode",
+              "lora_prompt", "custom_prompt"):
+        info.add_text(k, str(meta.get(k, "")))
+    info.add_text("loras", json.dumps(meta.get("loras") or []))
+    return info
+
+
+def save_with_metadata(image: Image.Image, meta: dict[str, Any],
+                       path: str | None = None) -> str:
+    """Save `image` as PNG with embedded generation metadata.
+
+    If `path` is given, write there (used by bulk-process to keep predictable
+    filenames inside its work directory). Otherwise allocate a temp PNG.
+    """
+    if path is None:
+        tmp = tempfile.NamedTemporaryFile(suffix=".png", delete=False, prefix="flux2_klein_")
+        tmp.close()
+        path = tmp.name
+    image.save(path, format="PNG", pnginfo=build_pnginfo(meta))
+    return path
+
+# ── Send arbitrary PIL → base / reference (used by the Depth/Pose tab) ──────
+
+def push_pil_to_base(img):
+    if img is None:
+        raise gr.Error("Nothing to send — generate it first.")
+    gr.Info("Sent to Base Image.")
+    return img
+
+
+def push_pil_to_reference(img, current_gallery):
+    if img is None:
+        raise gr.Error("Nothing to send — generate it first.")
+    current = list(current_gallery or [])
+    current.append(img)
+    gr.Info("Added to Reference Images.")
+    return current

lora_registry.py

@@ -0,0 +1,285 @@
+"""LoRA catalog + per-session custom LoRA loader.
+
+Two layers of state:
+
+* LORA_STYLES (module-level constant) — the curated catalog, identical for
+  every visitor.
+* dynamic_loras (gr.State, per session) — a session-private dict of any
+  custom HuggingFace LoRAs the user has added via the UI. Other sessions
+  never see them.
+
+Adapter loading itself is tracked in LOADED_ADAPTERS, which is a *legitimate*
+module-level set: it mirrors what's been attached to the shared pipe object
+in generation.py so we don't re-download/re-attach the same weights twice.
+"""
+from __future__ import annotations
+
+import uuid
+
+import gradio as gr
+
+from config import MAX_LORA_SLOTS
+
+
+FACE_SWAP_PROMPT = """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.
+FROM PICTURE 1 (strictly preserve):
+- Scene: lighting conditions, shadows, highlights, color temperature, environment, background
+- Head positioning: exact rotation angle, tilt, direction the head is facing
+- Expression: facial expression, micro-expressions, eye gaze direction, mouth position, emotion
+FROM PICTURE 2 (strictly preserve identity):
+- Facial structure: face shape, bone structure, jawline, chin
+- All facial features: eye color, eye shape, nose structure, lip shape and fullness, eyebrows
+- Hair: color, style, texture, hairline
+- Skin: texture, tone, complexion
+The replaced head must seamlessly match Picture 1's lighting and expression while maintaining the complete identity from Picture 2. High quality, photorealistic, sharp details, 4k."""
+
+# MAX number of simultaneous LoRA weight sliders to pre-render in the UI
+MAX_LORA_SLOTS = 6
+
+LORA_STYLES = [
+    {
+        "image": "https://huggingface.co/spaces/prithivMLmods/FLUX.2-Klein-LoRA-Studio/resolve/main/examples/image.webp",
+        "title": "None",
+        "adapter_name": None,
+        "repo": None,
+        "weights": None,
+        "default_prompt": None,
+        "default_weight": 1.0,
+    },
+    {
+        "title": "Klein-Delight-Style",
+        "adapter_name": "klein-delight",
+        "repo": "linoyts/Flux2-Klein-Delight-LoRA",
+        "weights": "pytorch_lora_weights.safetensors",
+        "default_prompt": "Relight the image to remove all existing lighting conditions and replace them with neutral, uniform illumination. Apply soft, evenly distributed lighting with no directional shadows, no harsh highlights, and no dramatic contrast. Maintain the original identity of all subjects exactly—preserve facial structure, skin tone, proportions, expressions, hair, clothing, and textures. Do not alter pose, camera angle, background geometry, or image composition. Lighting should appear balanced, and studio-neutral, similar to diffuse overcast or a soft lightbox setup. Ensure consistent exposure across the entire image with realistic depth and subtle shading only where necessary for form.",
+        "default_weight": 1.0,
+    },
+    {
+        "title": "Klein-Consistency",
+        "adapter_name": "klein-consistency",
+        "repo": "dx8152/Flux2-Klein-9B-Consistency",
+        "weights": "Klein-consistency.safetensors",
+        "default_prompt": None,
+        "default_weight": 0.3,
+    },
+    {
+        "title": "Best-Face-Swap",
+        "adapter_name": "face-swap",
+        "repo": "Alissonerdx/BFS-Best-Face-Swap",
+        "weights": "bfs_head_v1_flux-klein_9b_step3750_rank64.safetensors",
+        "default_prompt": FACE_SWAP_PROMPT,
+        "default_weight": 1.0,
+    },
+    {
+        "title": "NSFW v2",
+        "adapter_name": "nsfw-v2",
+        "repo": "diroverflo/FLux_Klein_9B_NSFW",
+        "weights": "Flux Klein - NSFW v2.safetensors",
+        "default_prompt": None,
+        "default_weight": 1.0,
+    },
+    {
+        "title": "Ultimate Upscaler Klein-9b",
+        "adapter_name": "Ultimate Upscaler",
+        "repo": "loras",
+        "weights": "Flux2-Klein-Image-RestoreV1.safetensors",
+        "default_prompt": "restore the image quality, remove any compression artefacts, remove any haze and soft edges, enrich the original with new intricate detail in all textures and surfaces creating a professional photorealistic photograph with natural lighting and skin texture.",
+        "default_weight": 1.0,
+    },    
+    {
+        "title": "High Resolution",
+        "adapter_name": "High Resolution",
+        "repo": "loras",
+        "weights": "HighResolution9B.safetensors",
+        "default_prompt": "High Resolution",
+        "default_weight": 1.0,
+    },       
+    {
+        "title": "Female Asshole",
+        "adapter_name": "Female Asshole",
+        "repo": "loras",
+        "weights": "femaleasshole-f2-klein-9b.safetensors",
+        "default_prompt": None,
+        "default_weight": 1.0,
+    },        
+    {
+        "title": "Realistic Nudes",
+        "adapter_name": "Realistic Nudes",
+        "repo": "loras",
+        "weights": "realistic_nudes_klein_v3.safetensors",
+        "default_prompt": None,
+        "default_weight": 1.0,
+    },            
+    {
+        "title": "Perky Pointy Puffy Breasts",
+        "adapter_name": "Perky Pointy Puffy Breasts",
+        "repo": "loras",
+        "weights": "PerkyPointyPuffy_v1.1_small_pointy_breasts_large_puffy_nipples.safetensors",
+        "default_prompt": "Small pointy breasts with large puffy nipples",
+        "default_weight": 1.0,
+    },       
+    {
+        "title": "Flat Chested",
+        "adapter_name": "Flat Chested",
+        "repo": "loras",
+        "weights": "Flux2-Klein-9b-FlatChested-v1.safetensors",
+        "default_prompt": "flat chested",
+        "default_weight": 1.5,
+    },           
+    {
+        "title": "Controllight",
+        "adapter_name": "Controllight",
+        "repo": "ControlLight/ControlLight",
+        "weights": "controllight.safetensors",
+        "default_prompt": None,
+        "default_weight": 1.0,
+    },      
+    {
+        "title": "RefControl - Depth",
+        "adapter_name": "RefConDep",
+        "repo": "thedeoxen/refcontrol-FLUX.2-klein-9B-reference-depth-lora",
+        "weights": "flux2_klein_9b_refcontrol_depth.safetensors",
+        "default_prompt": "refcontrol",
+        "default_weight": 1.0,
+    },         
+    {
+        "title": "RefControl - Pose",
+        "adapter_name": "RefConPos",
+        "repo": "thedeoxen/refcontrol-FLUX.2-klein-9B-reference-pose-lora",
+        "weights": "refcontrol_v2_poses.safetensors",
+        "default_prompt": "apply pose from image 1 with reference from image 2",
+        "default_weight": 1.0,
+    },    
+]
+
+
+
+# LOADED_ADAPTERS is the only piece of LoRA state that's legitimately global:
+# it just tracks which adapter names have been loaded onto the shared `pipe`
+# at least once, so we don't re-download/re-attach weights on every call.
+# It is NOT used to decide what a user can select — that comes from each
+# session's own `dynamic_loras_state` (a gr.State dict), so one user's custom
+# HuggingFace LoRA never shows up in another user's checklist.
+LOADED_ADAPTERS: set[str] = set()
+
+
+# ── Catalog lookup helpers ────────────────────────────────────────────────────
+
+def get_all_styles(dynamic_loras):
+    all_styles = list(LORA_STYLES)
+    for dynamic_lora in (dynamic_loras or {}).values():
+        all_styles.append(dynamic_lora)
+    return all_styles
+
+
+def get_selectable_styles(dynamic_loras):
+    """Return all styles except 'None' for the checkbox selector."""
+    return [s for s in get_all_styles(dynamic_loras) if s["adapter_name"] is not None]
+
+
+def get_style_by_title(title, dynamic_loras):
+    for style in get_all_styles(dynamic_loras):
+        if style["title"] == title:
+            return style
+    return None
+
+
+def get_style_by_adapter_name(adapter_name, dynamic_loras):
+    for style in get_all_styles(dynamic_loras):
+        if style["adapter_name"] == adapter_name:
+            return style
+    return None
+
+
+# ── LoRA selection → update weight sliders + LoRA prompt display ──────────────
+
+def update_weight_sliders(selected_titles, dynamic_loras):
+    """Given a list of selected LoRA titles, return updated visibility/label/value
+    for each of the MAX_LORA_SLOTS pre-rendered sliders, plus an update for the
+    LoRA prompt display box (visible when any selected LoRA has a default prompt).
+
+    The user's own prompt is never touched here — LoRA prompts are shown
+    separately and appended at inference time.
+    """
+    selected_styles = [
+        get_style_by_title(t, dynamic_loras)
+        for t in (selected_titles or [])
+        if get_style_by_title(t, dynamic_loras)
+    ]
+    slider_updates = []
+    for i in range(MAX_LORA_SLOTS):
+        if i < len(selected_styles):
+            style = selected_styles[i]
+            slider_updates.append(gr.update(
+                visible=True,
+                label=f"{style['title']} — weight",
+                value=style.get("default_weight", 1.0),
+            ))
+        else:
+            slider_updates.append(gr.update(visible=False, value=1.0))
+
+    # Collect default prompts from ALL selected LoRAs (not just one)
+    lora_prompts = [s.get("default_prompt") for s in selected_styles if s.get("default_prompt")]
+    if lora_prompts:
+        combined = "\n\n".join(lora_prompts)
+        lora_prompt_update = gr.update(value=combined, visible=True)
+    else:
+        lora_prompt_update = gr.update(value="", visible=False)
+
+    return slider_updates + [lora_prompt_update]
+
+
+# ── Dynamic HF LoRA loader (per-session) ─────────────────────────────────────
+
+def add_custom_lora(repo_id, weight_name, adapter_name, dynamic_loras_state):
+    """Adds a custom LoRA to *this session's* state only — other users never see
+    it in their selector. The adapter_name attached to the shared pipe is
+    still made globally unique (via a short random suffix) since the pipe
+    object itself is one shared resource across all sessions.
+    """
+    dynamic_loras = dict(dynamic_loras_state or {})
+    if not repo_id or not repo_id.strip():
+        return "Please enter a valid HuggingFace repo ID.", gr.update(), dynamic_loras
+    repo_id = repo_id.strip()
+    requested_name = adapter_name.strip() if adapter_name and adapter_name.strip() else None
+    try:
+        from huggingface_hub import model_info
+        info = model_info(repo_id)
+        actual_weight = weight_name.strip() if weight_name and weight_name.strip() else None
+        if not actual_weight:
+            for name in ["pytorch_lora_weights.safetensors", "lora.safetensors", "adapter_model.safetensors"]:
+                if any(f.filename == name for f in info.siblings):
+                    actual_weight = name
+                    break
+        if not actual_weight:
+            available = [f.filename for f in info.siblings if f.filename.endswith(('.safetensors', '.bin'))]
+            return (
+                f"No weight found. Available: {', '.join(available) or 'None'}",
+                gr.update(),
+                dynamic_loras,
+            )
+
+        base_name = (
+            "".join(c if c.isalnum() or c in "-_" else "_" for c in requested_name)
+            if requested_name else "custom"
+        )
+
+        static_names = {s["adapter_name"] for s in LORA_STYLES if s["adapter_name"]}
+        final_adapter_name = f"{base_name}_{uuid.uuid4().hex[:6]}"
+        while final_adapter_name in static_names or final_adapter_name in LOADED_ADAPTERS:
+            final_adapter_name = f"{base_name}_{uuid.uuid4().hex[:6]}"
+
+        custom_style = {
+            "image": "https://huggingface.co/spaces/prithivMLmods/FLUX.2-Klein-LoRA-Studio/resolve/main/examples/image.webp",
+            "title": f"Custom: {base_name}",
+            "adapter_name": final_adapter_name,
+            "repo": repo_id,
+            "weights": actual_weight,
+            "default_prompt": None,
+            "default_weight": 1.0,
+        }
+        dynamic_loras[final_adapter_name] = custom_style
+        new_choices = [s["title"] for s in get_selectable_styles(dynamic_loras)]
+        return f"✅ Added: {base_name} from {repo_id}", gr.update(choices=new_choices), dynamic_loras
+    except Exception as e:
+        return f"❌ Failed: {str(e)}", gr.update(), dynamic_loras

requirements.txt

@@ -1,12 +1,11 @@
 git+https://github.com/huggingface/diffusers.git
-transformers==4.57.3
 huggingface_hub
 sentencepiece
 torch==2.8.0
 bitsandbytes
 torchvision
 accelerate
-kernels # @vllm-flash-attn3
+kernels<0.10
 torchao
 spaces
 hf_xet
@@ -15,4 +14,6 @@ numpy
 peft
 av
 spandrel
-pyarrow 
+pillow-heif
+pyarrow 
+controlnet-aux

ui_theme.py

@@ -0,0 +1,64 @@
+"""Custom Gradio theme — extracted verbatim from the original app.py."""
+from typing import Iterable
+
+from gradio.themes import Soft
+from gradio.themes.utils import colors, fonts, sizes
+
+
+# Register a custom colour palette on the gradio colours module so it can be
+# referenced by name elsewhere (matching the original code).
+colors.orange_red = colors.Color(
+    name="orange_red",
+    c50="#FFF0E5", c100="#FFE0CC", c200="#FFC299", c300="#FFA366",
+    c400="#FF8533", c500="#FF4500", c600="#E63E00", c700="#CC3700",
+    c800="#B33000", c900="#992900", c950="#802200",
+)
+
+
+class OrangeRedTheme(Soft):
+    def __init__(
+        self,
+        *,
+        primary_hue: colors.Color | str = colors.gray,
+        secondary_hue: colors.Color | str = colors.orange_red,
+        neutral_hue: colors.Color | str = colors.slate,
+        text_size: sizes.Size | str = sizes.text_lg,
+        font: fonts.Font | str | Iterable[fonts.Font | str] = (
+            fonts.GoogleFont("Outfit"), "Arial", "sans-serif",
+        ),
+        font_mono: fonts.Font | str | Iterable[fonts.Font | str] = (
+            fonts.GoogleFont("IBM Plex Mono"), "ui-monospace", "monospace",
+        ),
+    ):
+        super().__init__(
+            primary_hue=primary_hue,
+            secondary_hue=secondary_hue,
+            neutral_hue=neutral_hue,
+            text_size=text_size,
+            font=font,
+            font_mono=font_mono,
+        )
+        super().set(
+            background_fill_primary="*primary_50",
+            background_fill_primary_dark="*primary_900",
+            body_background_fill="linear-gradient(135deg, *primary_200, *primary_100)",
+            body_background_fill_dark="linear-gradient(135deg, *primary_900, *primary_800)",
+            button_primary_text_color="white",
+            button_primary_text_color_hover="white",
+            button_primary_background_fill="linear-gradient(90deg, *secondary_500, *secondary_600)",
+            button_primary_background_fill_hover="linear-gradient(90deg, *secondary_600, *secondary_700)",
+            button_primary_background_fill_dark="linear-gradient(90deg, *secondary_600, *secondary_700)",
+            button_primary_background_fill_hover_dark="linear-gradient(90deg, *secondary_500, *secondary_600)",
+            slider_color="*secondary_500",
+            slider_color_dark="*secondary_600",
+            block_title_text_weight="600",
+            block_border_width="3px",
+            block_shadow="*shadow_drop_lg",
+            button_primary_shadow="*shadow_drop_lg",
+            button_large_padding="11px",
+            color_accent_soft="*primary_100",
+            block_label_background_fill="*primary_200",
+        )
+
+
+orange_red_theme = OrangeRedTheme()

upscale.py

@@ -0,0 +1,136 @@
+"""ESRGAN-family upscaler registry + tiled inference.
+
+Spandrel handles any ESRGAN-family .pth file — to add a model, just append an
+entry to UPSCALE_MODELS with its download URL and integer scale factor.
+"""
+from __future__ import annotations
+
+import gc
+import os
+import urllib.request
+
+import numpy as np
+import torch
+import gradio as gr
+from PIL import Image
+
+
+# ── Upscaler model registry ───────────────────────────────────────────────────
+# Add new entries here to extend the dropdown — spandrel handles any ESRGAN-
+# family .pth file automatically. scale= is the integer output multiplier.
+UPSCALE_MODELS = {
+    "None": {
+        "scale": None, "file": None, "url": None,
+    },
+    "2× — RealESRGAN (balanced)": {
+        "scale": 2,
+        "file": "RealESRGAN_x2plus.pth",
+        "url": "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.1/RealESRGAN_x2plus.pth",
+    },
+    "4× — RealESRGAN (balanced)": {
+        "scale": 4,
+        "file": "RealESRGAN_x4plus.pth",
+        "url": "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth",
+    },
+    "4× — UltraSharp (crisp)": {
+        "scale": 4,
+        "file": "4x-UltraSharpV2.pth",
+        "url": "https://huggingface.co/Kim2091/UltraSharpV2/resolve/main/4x-UltraSharpV2.pth",
+    },
+    "4× — Remacri (natural)": {
+        "scale": 4,
+        "file": "4x_foolhardy_Remacri.pth",
+        "url": "https://huggingface.co/FacehugmanIII/4x_foolhardy_Remacri/resolve/main/4x_foolhardy_Remacri.pth",
+    },
+    "4× — Nomos2 HQ DAT2 (Photography)": {
+        "scale": 4,
+        "file": "4xNomos2_hq_dat2.pth",
+        "url": "https://github.com/Phhofm/models/releases/download/4xNomos2_hq_dat2/4xNomos2_hq_dat2.pth",
+    },
+}
+
+
+def _upscale_tiled(
+    model_fn,
+    img_t: torch.Tensor,
+    tile: int = 512,
+    overlap: int = 32,
+) -> torch.Tensor:
+    """Run the upscaler model patch-by-patch to cap peak VRAM usage.
+
+    Overlapping tiles are averaged so there are no seam artifacts. Output is
+    assembled on CPU so only one tile lives on GPU at a time.
+    """
+    _, c, h, w = img_t.shape
+
+    # Probe output scale with a tiny crop rather than hard-coding it
+    with torch.no_grad():
+        probe = model_fn(img_t[:, :, :min(4, h), :min(4, w)])
+    scale = probe.shape[-1] // min(4, w)
+    del probe
+    torch.cuda.empty_cache()
+
+    out_h, out_w = h * scale, w * scale
+    canvas = torch.zeros(1, c, out_h, out_w, dtype=torch.float32)
+    weights = torch.zeros(1, 1, out_h, out_w, dtype=torch.float32)
+
+    step = max(tile - overlap, 1)
+    # Ensure the last tile always reaches the edge
+    ys = list(range(0, max(h - tile, 0), step)) + [max(h - tile, 0)]
+    xs = list(range(0, max(w - tile, 0), step)) + [max(w - tile, 0)]
+    ys = sorted(set(ys))
+    xs = sorted(set(xs))
+
+    for y0 in ys:
+        for x0 in xs:
+            y1 = min(y0 + tile, h)
+            x1 = min(x0 + tile, w)
+            patch = img_t[:, :, y0:y1, x0:x1]
+            with torch.no_grad():
+                out = model_fn(patch).cpu().float()
+            oy0, ox0 = y0 * scale, x0 * scale
+            oy1, ox1 = y1 * scale, x1 * scale
+            canvas[:, :, oy0:oy1, ox0:ox1] += out
+            weights[:, :, oy0:oy1, ox0:ox1] += 1.0
+
+    return (canvas / weights.clamp(min=1)).clamp(0, 1)
+
+
+def apply_realesrgan(image: Image.Image, model_key: str, device: torch.device) -> Image.Image:
+    """Upscale a PIL image using the model selected in UPSCALE_MODELS.
+
+    Weights are downloaded once to /tmp/realesrgan_weights/ and reused.
+    """
+    cfg = UPSCALE_MODELS[model_key]
+
+    try:
+        from spandrel import ImageModelDescriptor, ModelLoader
+    except ImportError:
+        raise gr.Error("spandrel is not installed. Add 'spandrel' to requirements.txt.")
+
+    cache_dir = "/tmp/realesrgan_weights"
+    os.makedirs(cache_dir, exist_ok=True)
+    cache_path = os.path.join(cache_dir, cfg["file"])
+
+    if not os.path.exists(cache_path):
+        print(f"Downloading {cfg['file']}…")
+        urllib.request.urlretrieve(cfg["url"], cache_path)
+
+    model = ModelLoader().load_from_file(cache_path)
+    if not isinstance(model, ImageModelDescriptor):
+        raise gr.Error(f"Loaded model is not a single-image descriptor: {type(model)}")
+
+    sr_model = model.model.to(device).eval()
+
+    img_np = np.array(image).astype(np.float32) / 255.0
+    img_t = torch.from_numpy(img_np).permute(2, 0, 1).unsqueeze(0).to(device)
+
+    out_t = _upscale_tiled(sr_model, img_t, tile=512, overlap=32)
+
+    # Free upscaler weights immediately — FLUX pipeline stays resident
+    del sr_model, img_t
+    gc.collect()
+    torch.cuda.empty_cache()
+
+    out_np = out_t.squeeze(0).permute(1, 2, 0).numpy()
+    return Image.fromarray((out_np * 255).astype(np.uint8))