app.py

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
from PIL import Image

from logging_utils import log_inference

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":
    from diffusers import Flux2KleinKVPipeline as _PipeClass
    _MODEL_REPO = "black-forest-labs/FLUX.2-klein-9b-kv"
else:
    from diffusers import Flux2KleinPipeline as _PipeClass
    _MODEL_REPO = "black-forest-labs/FLUX.2-klein-9B"
    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

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.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",
            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()
MAX_SEED = np.iinfo(np.int32).max

# ── 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"},
}

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_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 = set()


def get_all_styles(dynamic_loras):
    return list(LORA_STYLES) + list((dynamic_loras or {}).values())

def get_selectable_styles(dynamic_loras):
    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 s in get_all_styles(dynamic_loras):
        if s["title"] == title:
            return s
    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}")


# ── UI helper callbacks ──────────────────────────────────────────────────────

def update_weight_sliders(selected_titles, dynamic_loras):
    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):
            s = selected[i]
            updates.append(gr.update(visible=True,
                                     label=f"{s['title']} — weight",
                                     value=s.get("default_weight", 1.0)))
        else:
            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)]


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)


def on_fit_mode_change(fit_mode):
    """Pad colour swatch only relevant for Pad (color)."""
    return gr.update(visible=(fit_mode == "Pad (color)"))


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()
    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)
    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)
            with torch.no_grad():
                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, 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)
    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))


# ── 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(
        target=log_inference,
        args=(pil_images, result_image, prompt, seed, steps, guidance_scale,
              width, height, duration, success, error),
        kwargs={"lora_titles": lora_titles, "lora_weights": lora_weights,
                "upscale_factor": upscale_factor, "lora_prompt_text": lora_prompt_text},
        daemon=True,
    ).start()


# ── 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,
    canvas_mode, custom_width, custom_height,
    canvas_fit_mode, pad_color, dynamic_loras,
    *slider_values, progress=gr.Progress(track_tqdm=True),
):
    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.")
        if len(pil_images) > 2:
            gr.Warning("Face Swap uses only the Base image and the first Reference image.")
            pil_images = pil_images[:2]

    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:
        pipe.disable_lora()
    else:
        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=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])

    full_prompt = "\n".join(p for p in [
        (prompt or "").strip(),
        (lora_prompt_text or "").strip(),
        (custom_prompt_text or "").strip(),
    ] if p)

    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 = [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:
        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":
            kwargs["guidance_scale"] = guidance_scale
        image = pipe(**kwargs).images[0]
    except Exception as e:
        raise gr.Error(f"Inference failed: {e}")

    if upscale_factor and upscale_factor != "None":
        gc.collect(); torch.cuda.synchronize(); torch.cuda.empty_cache()
        try:
            image = apply_realesrgan(image, upscale_factor)
        except Exception as e:
            gr.Warning(f"Upscaling failed, returning {width}×{height} result: {e}")

    gc.collect(); torch.cuda.empty_cache()
    return image, seed, width, height


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):
    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 = 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,
            "repo": repo_id, "weights": actual_weight,
            "default_prompt": None, "default_weight": 1.0,
        }
        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 (unchanged) ────────────────────────────────────────

def add_custom_prompt(name, text, prompts_state, 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}"
    if name in prompts:
        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, 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]
    choices = list(prompts.keys())
    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:
        return gr.update(value="", visible=False)
    texts = [prompts_state[n] for n in selected_names if n in prompts_state]
    if texts:
        return gr.update(value="\n\n".join(texts), visible=True)
    return gr.update(value="", visible=False)


# ── UI ───────────────────────────────────────────────────────────────────────

css = """
#col-container { margin: 0 auto; max-width: 980px; }
#main-title h1 { font-size: 2.4em !important; }
#reference_gallery .grid-wrap { min-height: 120px }
.lora-weight-row { background: var(--block-background-fill); border-radius: 8px; padding: 4px 12px; margin-bottom: 4px; }
"""

# 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({})
    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"
        _logging_badge = "🟢 On" if _logging_on else "🔴 Off"

        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 using [FLUX.2-Klein-{MODEL_VARIANT}]({_MODEL_REPO}). "
            f"**Model:** `{MODEL_VARIANT}` · **Logging:** {_logging_badge}"
        )

        with gr.Tabs() as main_tabs:
            # ── Generate tab ─────────────────────────────────────────────────
            with gr.Tab("🎨 Generate", id="tab_generate"):
                with gr.Row(equal_height=False):
                    with gr.Column(scale=1):
                        base_image = gr.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,
                            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")
                        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")
                            steps = gr.Slider(label="Steps", minimum=1, maximum=50, value=4, step=1)
                            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 ─────────────────────────────────────────
                    with gr.Column(scale=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 + 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",
                )
                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"):
                            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,
                            ))

                with gr.Accordion("➕ Load Custom LoRA from HuggingFace", open=False):
                    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_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)

                with gr.Accordion("📝 Custom Prompts", open=False):
                    gr.Markdown("Save reusable prompt snippets for this session.")
                    custom_prompt_selector = gr.CheckboxGroup(
                        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)
                    with gr.Row():
                        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_btn = gr.Button("🗑️ Delete", variant="secondary", scale=1)

            # ── Crop / Fix Image tab ─────────────────────────────────────────
            with gr.Tab("✂️ Crop / Fix Image", id="tab_editor"):
                gr.Markdown(
                    "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=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")

            # ── 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)

            # ── 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."
                )

                # 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,
        inputs=[lora_selector, dynamic_loras_state],
        outputs=weight_sliders + [lora_prompt_display],
    )

    add_lora_btn.click(
        fn=add_custom_lora,
        inputs=[lora_repo_id, lora_weight_name, lora_adapter_name, dynamic_loras_state],
        outputs=[lora_status, lora_selector, dynamic_loras_state],
    )

    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],
    )
    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],
    )

    # 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,
                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],
    )

    # ── 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], 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])

    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])

    # ── 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__":
    # 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)