app.py

import os
import gc
import csv
import time
import random
import uuid
import zipfile

import gradio as gr
import spaces
import torch
import numpy as np
from PIL import Image

# ── Local modules — single source of truth for each concern ─────────────────
from config import (
    MODEL_VARIANT,
    MODEL_REPO,
    MAX_SEED,
    MAX_LORA_SLOTS,
    PERSISTENT_LORA_CATALOG_PATH,
    UNCENSORED_TE_REPO,
    UNCENSORED_TE_FILE,
)
from ui_theme import orange_red_theme
from upscale import UPSCALE_MODELS, apply_realesrgan
from lora_registry import (
    LORA_STYLES,
    LOADED_ADAPTERS,
    get_selectable_styles,
    get_style_by_title,
    update_weight_sliders,
    add_custom_lora,
    save_session_lora_to_catalog,
    fill_catalog_save_form,
    remove_lora,
    removable_catalog_titles,
    refresh_catalog_ui,
)

from image_utils import (
    fix_orientation,
    compute_canvas_dimensions,
    fit_to_canvas,
    on_base_image_change,
    make_solid_base_image,
    collect_reference_images,
    reference_info_text,
    compact_reference_slots,
    move_base_down,
    move_ref1_up,
    move_ref1_down,
    move_ref2_up,
    move_ref2_down,
    move_ref3_up,
    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,
    extend_editor_canvas,
    render_extend_schematic,
    resolve_output_download_path,
    latest_gallery_download_path,
)
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,
)

DEFAULT_PROJECT_NAME = "f2klora"
MAX_PROJECT_NAME_LEN = 12


def sanitize_project_name(name: str | None) -> str:
    """Alphanumeric only, max 12 chars. Default f2klora."""
    raw = (name or "").strip()
    cleaned = "".join(c for c in raw if c.isalnum())
    cleaned = cleaned[:MAX_PROJECT_NAME_LEN]
    return cleaned or DEFAULT_PROJECT_NAME


def make_run_stamp() -> str:
    """yymmddhhmmss — no separators."""
    return time.strftime("%y%m%d%H%M%S")


def make_download_basename(project: str | None, stamp: str | None = None,
                           batch_index: int | None = None) -> str:
    stem = f"{sanitize_project_name(project)}{stamp or make_run_stamp()}"
    if batch_index is not None:
        stem = f"{stem}{int(batch_index):02d}"
    return stem


def _tmp_named(basename: str, ext: str) -> str:
    ext = ext if ext.startswith(".") else f".{ext}"
    return f"/tmp/{basename}{ext}"


def save_simple_image(image: Image.Image, basename: str | None = None) -> str:
    """Save image without any metadata for privacy."""
    path = _tmp_named(basename or f"gen{uuid.uuid4().hex[:8]}", ".png")
    image.save(path, format="PNG")
    return path


def _build_full_prompt(prompt, lora_prompt_text, custom_prompt_text) -> str:
    return "\n".join(
        p for p in [
            (prompt or "").strip(),
            (lora_prompt_text or "").strip(),
            (custom_prompt_text or "").strip(),
        ] if p
    )


def save_webp_from_image(image: Image.Image, basename: str | None = None) -> str:
    path = _tmp_named(basename or f"gen{uuid.uuid4().hex[:8]}", ".webp")
    image.convert("RGB").save(path, format="WEBP", quality=90, method=4)
    return path


def save_webp_from_path(path, basename: str | None = None) -> str | None:
    from image_utils import _gallery_item_path
    p = _gallery_item_path(path) or (path if isinstance(path, str) else None)
    if not p:
        return None
    try:
        img = Image.open(p).convert("RGB")
    except Exception:
        return None
    return save_webp_from_image(img, basename=basename)


def save_prompt_txt(text: str | None, basename: str | None = None) -> str | None:
    text = (text or "").strip()
    if not text:
        return None
    path = _tmp_named(basename or f"prompt{uuid.uuid4().hex[:8]}", ".txt")
    with open(path, "w", encoding="utf-8") as f:
        f.write(text)
        if not text.endswith("\n"):
            f.write("\n")
    return path


def copy_as_named_png(src_path, basename: str) -> str | None:
    """Copy an existing PNG to a project+timestamp name for download."""
    from image_utils import _gallery_item_path
    import shutil
    p = _gallery_item_path(src_path) or (src_path if isinstance(src_path, str) else None)
    if not p or not os.path.isfile(p):
        return None
    dest = _tmp_named(basename, ".png")
    if os.path.abspath(p) == os.path.abspath(dest):
        return p
    try:
        shutil.copy2(p, dest)
        return dest
    except Exception:
        return p


def resolve_download_bundle(selected_path, gallery_value, last_prompt_text, project_name):
    """PNG + WebP + prompt txt named project+yymmddhhmmss.(png|webp|txt)."""
    png_src = resolve_output_download_path(selected_path, gallery_value)
    base = make_download_basename(project_name)
    png = copy_as_named_png(png_src, base) if png_src else None
    webp = save_webp_from_path(png_src, basename=base) if png_src else None
    prompt_file = save_prompt_txt(last_prompt_text, basename=base)
    return png, webp, prompt_file


# ── Download tracking (warn before overwriting undownloaded outputs) ─────────

def _norm_img_path(item) -> str | None:
    from image_utils import _gallery_item_path
    p = _gallery_item_path(item)
    if p:
        return os.path.abspath(p)
    if isinstance(item, str) and item.strip():
        return os.path.abspath(item) if os.path.isabs(item) else item.strip()
    return None


def gallery_image_paths(gallery_value) -> list[str]:
    paths: list[str] = []
    seen: set[str] = set()
    for item in gallery_value or []:
        p = _norm_img_path(item)
        if p and p not in seen:
            seen.add(p)
            paths.append(p)
    return paths


def format_download_status(pending_list, gallery_value) -> str:
    paths = gallery_image_paths(gallery_value)
    pending = set(pending_list or [])
    undownloaded = [p for p in paths if p in pending]
    if not paths:
        return "*No generated images yet.*"
    if not undownloaded:
        return f"✅ All **{len(paths)}** gallery image(s) marked downloaded."
    names = ", ".join(f"`{os.path.basename(p)}`" for p in undownloaded[:4])
    extra = f" +{len(undownloaded) - 4} more" if len(undownloaded) > 4 else ""
    return (
        f"⚠️ **{len(undownloaded)}/{len(paths)}** not downloaded yet: {names}{extra}. "
        f"Use ⬇️ PNG/WebP or the gallery ↓ icon."
    )


def sync_download_tracking(gallery_value, pending_list, downloaded_list):
    """Keep pending in sync with gallery contents.

    New gallery paths not yet marked downloaded become pending.
    Paths that left the gallery drop out of pending.
    """
    pending = set(pending_list or [])
    downloaded = set(downloaded_list or [])
    paths = gallery_image_paths(gallery_value)
    path_set = set(paths)
    pending = {p for p in pending if p in path_set}
    for p in paths:
        if p not in downloaded:
            pending.add(p)
    pending_out = [p for p in paths if p in pending]  # stable order
    downloaded_out = sorted(downloaded)
    return pending_out, downloaded_out, format_download_status(pending_out, gallery_value)


def _mark_paths_downloaded(paths_to_mark, gallery_value, pending_list, downloaded_list):
    pending = set(pending_list or [])
    downloaded = set(downloaded_list or [])
    marked = []
    for raw in paths_to_mark or []:
        p = _norm_img_path(raw)
        if not p:
            # bare filename / URL fragment from gallery JS
            s = str(raw or "").strip()
            if not s:
                continue
            base = os.path.basename(s.split("?")[0].split("#")[0])
            for gp in gallery_image_paths(gallery_value):
                if os.path.basename(gp) == base or base in gp or gp.endswith(base):
                    p = gp
                    break
            if not p and base:
                # still record basename key so status can clear if paths match later
                p = base
        if not p:
            continue
        downloaded.add(p)
        pending.discard(p)
        # also clear any gallery path sharing basename
        base = os.path.basename(p)
        for gp in list(pending):
            if os.path.basename(gp) == base:
                pending.discard(gp)
                downloaded.add(gp)
        marked.append(p)
    paths = gallery_image_paths(gallery_value)
    pending_out = [p for p in paths if p in pending]
    return pending_out, sorted(downloaded), format_download_status(pending_out, gallery_value)


def mark_current_output_downloaded(selected_path, gallery_value, pending_list, downloaded_list):
    """Mark the selected (or latest) gallery image as downloaded."""
    p = resolve_output_download_path(selected_path, gallery_value)
    return _mark_paths_downloaded([p] if p else [], gallery_value, pending_list, downloaded_list)


def mark_from_gallery_signal(signal, gallery_value, pending_list, downloaded_list):
    """Mark download from gallery ↓ icon (JS writes URL/filename into signal)."""
    raw = (signal or "").strip()
    if not raw:
        return (
            list(pending_list or []),
            list(downloaded_list or []),
            format_download_status(pending_list, gallery_value),
            gr.update(value=""),
        )
    # JS appends "|timestamp" so repeated clicks still fire .change
    token = raw.split("|", 1)[0].strip()
    pending, downloaded, status = _mark_paths_downloaded(
        [token], gallery_value, pending_list, downloaded_list,
    )
    return pending, downloaded, status, gr.update(value="")


def warn_undownloaded_before_generate(pending_list, gallery_value):
    """Disabled for now — undownloaded tracking will return later.

    Still clears selected_output_state when chained before generate.
    """
    return None
    
# ── Model load ──────────────────────────────────────────────────────────────
# Pipeline class depends on MODEL_VARIANT and is the only thing here that
# can't live in config.py (config must stay torch/diffusers-free).
if MODEL_VARIANT == "9B-KV":
    from diffusers import Flux2KleinKVPipeline as _PipeClass
else:
    from diffusers import Flux2KleinPipeline as _PipeClass

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

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}")

# ── Replace text encoder with abliterated (uncensored) version ───────────
try:
    from huggingface_hub import hf_hub_download
    from safetensors.torch import load_file
    print(f"Downloading abliterated text encoder from {UNCENSORED_TE_REPO}...")
    te_path = hf_hub_download(repo_id=UNCENSORED_TE_REPO, filename=UNCENSORED_TE_FILE)
    print(f"Loading abliterated weights from {te_path}...")
    state_dict = load_file(te_path)
    pipe.text_encoder.load_state_dict(state_dict, strict=True)
    pipe.text_encoder.to(device=device, dtype=torch.bfloat16)
    print("Abliterated text encoder loaded — safety filters removed.")
except Exception as e:
    print(f"Abliterated text encoder unavailable ({e}) — using stock encoder.")


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

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


# ── Logging (disabled) ───────────────────────────────────────────────────────

def _spawn_log(*_args, **_kwargs):
    """No-op — logging intentionally disabled."""
    return


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


# ── Single / batch infer (generator → streams into gr.Gallery) ───────────────

def infer(
    base_image, ref1, ref2, ref3, 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, project_name,
    dynamic_loras, *slider_values, progress=gr.Progress(track_tqdm=True),
):
    """Generator. Streams a list of PNG paths into the output 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.")
    reference_images = collect_reference_images(ref1, ref2, ref3)
    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))
    project = sanitize_project_name(project_name)

    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]

    results = []
    last_seed_text = ""
    full_prompt = _build_full_prompt(prompt, lora_prompt_text, custom_prompt_text)
    # One stamp for the whole batch; multi-run items get 00/01/... suffix (no hyphen).
    run_stamp = make_run_stamp()
    prompt_base = make_download_basename(project, run_stamp)
    prompt_file = save_prompt_txt(full_prompt, basename=prompt_base)
    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]
        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,
            )
            item_base = make_download_basename(
                project, run_stamp,
                batch_index=i if batch_count > 1 else None,
            )
            png_path = save_simple_image(image, basename=item_base)
            webp_path = save_webp_from_image(image, basename=item_base)
            results.append(png_path)
            last_seed_text = str(used_seed)
            # Yield download paths explicitly. Relying only on output_gallery.change
            # fails on the first generate in a virgin session (buttons stay empty
            # until some later interaction re-triggers the change chain).
            yield results, last_seed_text, png_path, webp_path, prompt_file, full_prompt
        except Exception as e:
            png = latest_gallery_download_path(results)
            fail_base = make_download_basename(project)
            webp = save_webp_from_path(png, basename=fail_base) if png else None
            png_named = copy_as_named_png(png, fail_base) if png else None
            yield (
                results,
                f"Batch {i+1}/{batch_count} failed: {e}",
                png_named, webp, prompt_file, full_prompt,
            )
            

# ── 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."""
    if not input_files:
        raise gr.Error("Upload at least one image first.")

    work_dir = _new_bulk_workdir()
    results = []
    succeeded = 0
    total = len(input_files)

    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]

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

            stem = os.path.splitext(fname)[0]
            out_path = os.path.join(work_dir, f"{i:03d}_{stem}.png")
            image.save(out_path, format="PNG")   # plain save, no metadata
            results.append(out_path)
            succeeded += 1
        except Exception as e:
            print(f"Bulk image {i+1} failed: {e}")

    return results

# ── Custom prompt manager (session-local) ────────────────────────────────────

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

# Shared viewport height so Base / Reference / Output feel the same size.
# Keep this moderate — oversized Gallery CSS previously split the reference
# panel into a huge empty pane + tiny control strip.
_IMAGE_BOX_H = 320
_OUTPUT_GALLERY_H = 520

css = f"""
#col-container {{ margin: 0 auto; max-width: 1100px; }}
#main-title h1 {{ font-size: 2.4em !important; }}
.lora-weight-row {{ background: var(--block-background-fill); border-radius: 8px; padding: 4px 12px; margin-bottom: 4px; }}
#used_seed textarea {{ min-height: 0 !important; height: 2.2rem !important; }}
/* Output gallery: avoid huge empty preview chrome / forced scrollbars */
#output_gallery {{ min-height: {_OUTPUT_GALLERY_H}px; }}
#output_gallery .grid-wrap,
#output_gallery .gallery-container,
#output_gallery .thumbnail-item,
#output_gallery .preview-image,
#output_gallery img {{
  max-height: {_OUTPUT_GALLERY_H - 48}px !important;
  object-fit: contain !important;
}}
#output_gallery .preview {{
  max-height: {_OUTPUT_GALLERY_H - 24}px !important;
  overflow: hidden !important;
}}
.slot-move-row button {{ min-width: 2.4rem !important; }}
"""

# Enter inserts newline in multi-line textboxes (Gradio default often submits).
# Shift+Enter also inserts newline for muscle-memory parity with chat UIs.
_TEXTBOX_NEWLINE_JS = """
() => {
  const isMulti = (el) => {
    if (!el || el.tagName !== 'TEXTAREA') return false;
    if (el.closest('#used_seed') || el.closest('#project_name')) return false;
    return true;
  };
  const onKey = (e) => {
    if (e.key !== 'Enter' || e.isComposing) return;
    const t = e.target;
    if (!isMulti(t)) return;
    // Always keep newline behaviour; never submit the form from a prompt box.
    e.stopPropagation();
    // Browser already inserts newline on plain Enter in textarea;
    // for Shift+Enter some hosts swallow it — insert manually if needed.
    if (e.shiftKey) {
      e.preventDefault();
      const start = t.selectionStart ?? t.value.length;
      const end = t.selectionEnd ?? start;
      const v = t.value;
      t.value = v.slice(0, start) + '\\n' + v.slice(end);
      const pos = start + 1;
      t.selectionStart = t.selectionEnd = pos;
      t.dispatchEvent(new Event('input', { bubbles: true }));
    }
  };
  document.addEventListener('keydown', onKey, true);
}
"""

# 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)
    last_prompt_state = gr.State("")
    # Paths still in the gallery that have not been marked downloaded this session.
    pending_download_state = gr.State([])
    downloaded_images_state = gr.State([])
    # Filled by JS when the gallery's built-in ↓ icon is clicked.
    gallery_dl_signal = gr.Textbox(
        value="", visible=False, elem_id="gallery_dl_signal",
    )
    # Remember per-title LoRA weights / prompts across add/remove so existing
    # values don't snap back to catalog defaults when the selection changes.
    lora_weight_memory_state = gr.State({})
    lora_prompt_memory_state = gr.State({})
    lora_prev_selected_state = gr.State([])

    with gr.Column(elem_id="col-container"):
        gr.Markdown("# **flux2klein lora playground**", 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}`"
        )

        with gr.Tabs() as main_tabs:
            # ── Generate tab ─────────────────────────────────────────────────
            with gr.Tab("🎨 Generate", id="tab_generate"):
                # equal_height=False: otherwise the short seed box stretches to match
                # the tall left column (base + reference).
                with gr.Row(equal_height=False):
                    with gr.Column(scale=1):
                        base_image = gr.Image(
                            label="Base Image", type="pil",
                            sources=["upload", "clipboard"],
                            height=_IMAGE_BOX_H,
                            elem_id="base_image",
                        )
                        with gr.Row(elem_classes="slot-move-row"):
                            base_down_btn = gr.Button("↓ Base → Ref1", size="sm")
                        # t2i-style workflows still need a base image; solid colour
                        # often becomes the background for i2i LoRAs used as t2i.
                        with gr.Row():
                            base_solid_color = gr.ColorPicker(
                                label="Solid base colour",
                                value="#FFFFFF",
                                scale=1,
                                elem_id="base_solid_color",
                            )
                            make_solid_base_btn = gr.Button(
                                "⬜ Use solid base",
                                size="sm",
                                scale=1,
                                elem_id="make_solid_base_btn",
                            )
                        gr.Markdown(
                            "*No photo? Use a solid base for t2i-style runs. "
                            "Size follows Custom canvas W×H if set, else 1024×1024.*",
                        )
                        size_info = gr.Markdown("*No image uploaded yet*")
                        run_button_top = gr.Button(
                            "▶ Generate", variant="primary", size="lg",
                            elem_id="run_button_top",
                        )

                        # Progressive single-image refs (max 3). Ref2 appears after
                        # Ref1 is set; Ref3 after Ref2. Deleting a middle slot packs
                        # remaining refs upward. Reorder with ↑/↓ (includes Base).
                        ref1 = gr.Image(
                            label="Reference 1 — optional", type="pil",
                            sources=["upload", "clipboard"],
                            height=_IMAGE_BOX_H,
                            elem_id="ref1",
                        )
                        with gr.Row(elem_classes="slot-move-row"):
                            ref1_up_btn = gr.Button("↑", size="sm", scale=0)
                            ref1_down_btn = gr.Button("↓", size="sm", scale=0)
                        ref2 = gr.Image(
                            label="Reference 2 — optional", type="pil",
                            sources=["upload", "clipboard"],
                            height=_IMAGE_BOX_H,
                            visible=False,
                            elem_id="ref2",
                        )
                        with gr.Row(elem_classes="slot-move-row"):
                            ref2_up_btn = gr.Button("↑", size="sm", scale=0, visible=False)
                            ref2_down_btn = gr.Button("↓", size="sm", scale=0, visible=False)
                        ref3 = gr.Image(
                            label="Reference 3 — optional", type="pil",
                            sources=["upload", "clipboard"],
                            height=_IMAGE_BOX_H,
                            visible=False,
                            elem_id="ref3",
                        )
                        with gr.Row(elem_classes="slot-move-row"):
                            ref3_up_btn = gr.Button("↑", size="sm", scale=0, visible=False)
                        reference_info = gr.Markdown("📷 No reference images")
                        gr.Markdown(
                            "*Up to 3 reference images. Next box appears after you fill the previous one. "
                            "Clearing a slot shifts the others up. Use ↑/↓ to reorder Base + refs. "
                            "Face Swap uses Base + Reference 1 only.*"
                        )

                        prompt = gr.Textbox(
                            label="Prompt",
                            lines=3,
                            max_lines=12,
                            placeholder="Describe the edit, or leave blank for style-only LoRAs. Enter = new line.",
                        )
                        lora_prompt_display = gr.Textbox(
                            label="LoRA prompts (auto-filled from selection — editable for this run)",
                            interactive=True,
                            visible=True,
                            lines=3,
                            max_lines=16,
                            value="",
                            placeholder="Tick LoRAs above to auto-fill. Edits are kept when you change selection.",
                            info="Filled when you tick LoRAs. Edit freely for the current generate; "
                                 "does not change the stored catalog default.",
                        )
                        custom_prompt_display = gr.Textbox(
                            label="Custom Prompts (auto-appended)",
                            interactive=False, visible=False, lines=3, max_lines=12,
                        )
                        run_button = gr.Button("▶ Generate", variant="primary", size="lg")

                    with gr.Column(scale=1):
                        output_gallery = gr.Gallery(
                            label="Output", type="filepath", columns=2, rows=1,
                            height=_OUTPUT_GALLERY_H,
                            allow_preview=True, preview=False,
                            object_fit="contain", show_label=True,
                            elem_id="output_gallery",
                        )
                        with gr.Row():
                            used_seed = gr.Textbox(
                                label="🌱 Seed used (last run)",
                                interactive=False,
                                lines=1,
                                max_lines=1,
                                elem_id="used_seed",
                                scale=2,
                            )
                            project_name = gr.Textbox(
                                label="Project short name",
                                value=DEFAULT_PROJECT_NAME,
                                max_lines=1,
                                lines=1,
                                max_length=MAX_PROJECT_NAME_LEN,
                                placeholder=DEFAULT_PROJECT_NAME,
                                info="Max 12 letters/digits. Used in download filenames.",
                                scale=1,
                                elem_id="project_name",
                            )
                        with gr.Row():
                            download_png_btn = gr.DownloadButton(
                                label="⬇️ PNG",
                                value=None,
                                variant="primary",
                                size="sm",
                                scale=1,
                                elem_id="download_png_btn",
                            )
                            download_webp_btn = gr.DownloadButton(
                                label="⬇️ WebP",
                                value=None,
                                variant="secondary",
                                size="sm",
                                scale=1,
                                elem_id="download_webp_btn",
                            )
                            download_prompt_btn = gr.DownloadButton(
                                label="⬇️ Prompt",
                                value=None,
                                variant="secondary",
                                size="sm",
                                scale=1,
                                elem_id="download_prompt_btn",
                            )
                        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")
                        download_status = gr.Markdown(
                            "*No generated images yet.*",
                            elem_id="download_status",
                        )
                        gr.Markdown(
                            "*Downloads: `project` + `yymmddhhmmss` + `.png/.webp/.txt` "
                            f"(default project `{DEFAULT_PROJECT_NAME}`). "
                            "Click a gallery thumbnail before Download / Send→; otherwise latest. "
                            "Prompt file = last run's combined user + LoRA + custom text. "
                            "PNG/WebP or the gallery ↓ icon marks an image downloaded; "
                            "Generate warns if undownloaded images would be replaced.*"
                        )

                with gr.Row():
                    gr.Markdown("### 🎨 Select LoRA(s)", elem_classes=["lora-heading"])
                    reload_catalog_btn = gr.Button(
                        "🔄 Reload catalog", size="sm", scale=0,
                    )
                lora_selector = gr.CheckboxGroup(
                    choices=[s["title"] for s in get_selectable_styles({})],
                    value=[], label="Active LoRAs — tick one or more",
                )
                catalog_load_status = gr.Markdown("", visible=True)
                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 (HF repo or local path)", open=False):
                    gr.Markdown(
                        "Add any FLUX.2-Klein-compatible LoRA from a **HuggingFace repo** "
                        "(`user/repo`) or a **local path** (file or directory), e.g. "
                        "`/loras-flux/my.safetensors` or `/loras-flux/foo/bar/male`. "
                        "**Import is always session-only** — try it first, then optionally "
                        f"save it to the catalog JSON (`{PERSISTENT_LORA_CATALOG_PATH}`). "
                        "Duplicates are blocked by title, repo+filename, and sha256 when available."
                    )
                    with gr.Row():
                        lora_repo_id = gr.Textbox(
                            label="HF repo ID or local path",
                            placeholder="user/repo  or  user/repo/sub/model.safetensors  or  /loras-flux/my.safetensors",
                            info="Nested HF paths OK: user/repo/folder/model.safetensors",
                        )
                    with gr.Row():
                        lora_weight_name = gr.Textbox(
                            label="Weight path inside repo (optional)",
                            placeholder="subfolder/model.safetensors",
                            info="Use for nested files if not included in the repo field.",
                        )
                        lora_adapter_name = gr.Textbox(label="Adapter name (optional)", placeholder="my-lora")
                    with gr.Row():
                        add_lora_btn = gr.Button("Add LoRA (session only)", variant="primary")
                        lora_status = gr.Textbox(label="Status", interactive=False)

                    gr.Markdown("#### 💾 Save tried LoRA to catalog")
                    gr.Markdown(
                        "After testing a session LoRA, save it here so it appears for everyone "
                        "on the next load. UI saves are **not** admin-approved; set "
                        "`admin_approved: true` in the JSON yourself. Set `active: false` to "
                        "archive/hide without deleting."
                    )
                    catalog_save_select = gr.Dropdown(
                        label="Session LoRA to save",
                        choices=[], value=None, interactive=True,
                    )
                    with gr.Row():
                        catalog_save_title = gr.Textbox(
                            label="Catalog title", placeholder="My LoRA name", scale=2,
                        )
                        catalog_save_weight = gr.Slider(
                            label="Default weight", minimum=0.0, maximum=2.0,
                            step=0.05, value=1.0, scale=1,
                        )
                    catalog_save_prompt = gr.Textbox(
                        label="Default prompt (optional)", lines=2,
                        placeholder="Safe ready-to-go prompt auto-appended when selected",
                    )
                    catalog_save_triggers = gr.Textbox(
                        label="Known triggers (optional)", lines=3,
                        placeholder=(
                            "One per line or comma-separated. Docs only — not auto-appended.\n"
                            "e.g. small penis, large penis, flaccid penis, erect penis"
                        ),
                        info="Can include mutually exclusive keywords; pick what you need in the prompt.",
                    )
                    catalog_save_notes = gr.Textbox(
                        label="Notes (optional)", lines=2,
                        placeholder="Usage notes, caveats, pairing tips…",
                    )
                    with gr.Row():
                        catalog_save_btn = gr.Button(
                            "💾 Save to catalog", variant="secondary",
                        )
                        catalog_save_status = gr.Textbox(label="Catalog status", interactive=False)

                    gr.Markdown("#### 🗑️ Remove LoRA")
                    gr.Markdown(
                        "Remove a **session** custom LoRA, or a catalog entry that is **not** "
                        "`admin_approved`. Admin-approved entries can only be archived via JSON "
                        "(`active: false`)."
                    )
                    with gr.Row():
                        catalog_remove_select = gr.Dropdown(
                            label="LoRA to remove",
                            choices=removable_catalog_titles(),
                            value=None, interactive=True, scale=2,
                        )
                        catalog_remove_btn = gr.Button("🗑️ Remove", variant="stop", scale=1)
                    catalog_remove_status = gr.Textbox(label="Remove 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)

                # Full-width advanced block at the bottom of Generate tab
                with gr.Accordion("⚙️ Advanced Settings", open=False):
                    with gr.Row():
                        with gr.Column(scale=1):
                            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",
                            )
                        with gr.Column(scale=1):
                            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."),
                    )
                    with gr.Row():
                        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,
                        )

                    with gr.Accordion("📋 Selected LoRA details (repo / triggers / notes)", open=False):
                        selected_lora_details = gr.Markdown(
                            value="*Tick one or more LoRAs above to see full repo paths, "
                                  "known triggers, and notes.*",
                            elem_id="selected_lora_details",
                        )

            # ── 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=False,
                    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")

                # ── Extend canvas section ────────────────────────────────────────────────
                # Uses the editor's current composite as the source so cropping + painting
                # happen first, then we grow the canvas around the result. Output is
                # loaded back into the same editor — Send → Base / Reference from there.
                with gr.Accordion("📐 Extend canvas (add padding around image)", open=True):
                    gr.Markdown(
                        "Grow the editor image's canvas by a percentage in any combination "
                        "of directions. Percentages are relative to the *current* image "
                        "size — `Down = 100` doubles the height with the image on top. "
                        "The result replaces the editor contents so you can crop again or "
                        "send it to Base / Reference with the buttons above."
                    )
                    with gr.Row():
                        ext_up    = gr.Number(label="Up %",    value=0,   minimum=0, precision=2)
                        ext_down  = gr.Number(label="Down %",  value=0,   minimum=0, precision=2)
                        ext_left  = gr.Number(label="Left %",  value=0,   minimum=0, precision=2)
                        ext_right = gr.Number(label="Right %", value=0,   minimum=0, precision=2)
                    with gr.Row():
                        ext_fill = gr.ColorPicker(label="Fill colour", value="#000000")
                        extend_btn = gr.Button("📐 Extend canvas", variant="primary")
                    with gr.Row():
                        ext_schematic = gr.Image(
                            label="Layout preview (red outline = current image)",
                            type="pil", interactive=False, height=220,
                        )
                        ext_info = gr.Markdown("*Upload something into the editor first.*")

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

                pose_source_state = gr.State(None)
                pose_keypoints_state = gr.State([])

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

    # Lightweight UI handlers: hide Gradio progress. On ZeroGPU/Spaces, the
    # default spinner often sticks on pure gr.update visibility changes
    # (LoRA weight sliders, canvas size text) until another event flushes UI.
    _ui = dict(show_progress="hidden")

    base_image.upload(fn=reencode_upload, inputs=[base_image], outputs=[base_image], **_ui)
    base_image.change(fn=on_base_image_change, inputs=[base_image], outputs=[size_info], **_ui)
    make_solid_base_btn.click(
        fn=make_solid_base_image,
        inputs=[base_solid_color, canvas_mode, custom_width, custom_height],
        outputs=[base_image],
        show_progress="hidden",
    ).then(
        fn=on_base_image_change, inputs=[base_image], outputs=[size_info], **_ui,
    )

    # Progressive ref slots: pack non-empty images upward on any change so
    # deleting Ref1 shifts Ref2/3 up instead of wiping them.
    for _ref in (ref1, ref2, ref3):
        _ref.upload(fn=reencode_upload, inputs=[_ref], outputs=[_ref], **_ui)

    _ref_compact_outputs = [
        ref1, ref2, ref3, reference_info,
        ref2_up_btn, ref2_down_btn, ref3_up_btn,
    ]
    for _ref in (ref1, ref2, ref3):
        _ref.change(
            fn=compact_reference_slots,
            inputs=[ref1, ref2, ref3],
            outputs=_ref_compact_outputs,
            **_ui,
        )

    _slot_inputs = [base_image, ref1, ref2, ref3]
    _move_outputs = [
        base_image, ref1, ref2, ref3, reference_info,
        ref2_up_btn, ref2_down_btn, ref3_up_btn,
    ]
    base_down_btn.click(fn=move_base_down, inputs=_slot_inputs, outputs=_move_outputs, **_ui)
    ref1_up_btn.click(fn=move_ref1_up, inputs=_slot_inputs, outputs=_move_outputs, **_ui)
    ref1_down_btn.click(fn=move_ref1_down, inputs=_slot_inputs, outputs=_move_outputs, **_ui)
    ref2_up_btn.click(fn=move_ref2_up, inputs=_slot_inputs, outputs=_move_outputs, **_ui)
    ref2_down_btn.click(fn=move_ref2_down, inputs=_slot_inputs, outputs=_move_outputs, **_ui)
    ref3_up_btn.click(fn=move_ref3_up, inputs=_slot_inputs, outputs=_move_outputs, **_ui)

    # Every browser open re-reads the bucket JSON and refreshes selector choices.
    # Without this, choices stay frozen at process start and newly saved LoRAs
    # (e.g. thickcum) are "already in catalog" but invisible in new sessions.
    def _on_page_load(dynamic_loras, selected):
        sel_upd, save_upd, rem_upd, dyn = refresh_catalog_ui(dynamic_loras, selected)
        # Count from live catalog after reload (gr.update is not always a plain dict).
        n = len(get_selectable_styles(dyn))
        status = f"*Catalog loaded — **{n}** active LoRA(s).*"
        return sel_upd, save_upd, rem_upd, dyn, status

    demo.load(
        fn=_on_page_load,
        inputs=[dynamic_loras_state, lora_selector],
        outputs=[lora_selector, catalog_save_select, catalog_remove_select,
                 dynamic_loras_state, catalog_load_status],
        show_progress="hidden",
    )
    # Must be registered inside the Blocks context (Gradio rejects load outside).
    demo.load(fn=None, js=_TEXTBOX_NEWLINE_JS)
    # Capture clicks on Gradio Gallery's built-in download (↓) control.
    demo.load(
        fn=None,
        js="""
() => {
  if (window.__fluxGalleryDlHook) return;
  window.__fluxGalleryDlHook = true;
  const setSignal = (val) => {
    const root = document.getElementById('gallery_dl_signal');
    if (!root) return;
    const ta = root.querySelector('textarea, input');
    if (!ta) return;
    ta.value = val || '';
    ta.dispatchEvent(new Event('input', { bubbles: true }));
  };
  document.addEventListener('click', (e) => {
    const t = e.target;
    if (!t || !t.closest) return;
    const gal = t.closest('#output_gallery');
    if (!gal) return;
    // Gradio download control: anchor with download attr, or button near download icon.
    const a = t.closest('a[download], a.download-link, a[href*="file="]');
    const btn = t.closest('button');
    let href = '';
    if (a && a.href) {
      href = a.getAttribute('download') || a.href;
    } else if (btn) {
      const label = (btn.getAttribute('aria-label') || btn.title || btn.textContent || '').toLowerCase();
      if (!(label.includes('download') || label.includes('save') || btn.innerHTML.includes('download'))) {
        // still allow if nested svg title looks like download
        const svgTitle = (btn.querySelector('title')?.textContent || '').toLowerCase();
        if (!svgTitle.includes('download') && !btn.querySelector('[data-testid*="download"]')) {
          return;
        }
      }
      const nearA = btn.closest('a') || btn.querySelector('a') || gal.querySelector('a[download]');
      href = (nearA && (nearA.getAttribute('download') || nearA.href)) || '';
      if (!href) {
        // fallback: selected/preview image src basename
        const img = gal.querySelector('.preview img, .thumbnail-lg img, img');
        href = (img && (img.currentSrc || img.src)) || 'gallery-download';
      }
    } else {
      return;
    }
    if (!href) return;
    try {
      const u = href.startsWith('http') || href.startsWith('blob:') || href.startsWith('/')
        ? href : href;
      const base = (u.split('/').pop() || u).split('?')[0];
      setSignal(base + '|' + Date.now());
    } catch (_) {
      setSignal(String(href) + '|' + Date.now());
    }
  }, true);
}
""",
    )
    reload_catalog_btn.click(
        fn=_on_page_load,
        inputs=[dynamic_loras_state, lora_selector],
        outputs=[lora_selector, catalog_save_select, catalog_remove_select,
                 dynamic_loras_state, catalog_load_status],
        show_progress="minimal",
    )

    # update_weight_sliders is the one imported from lora_registry now.
    # Only wire .change — also binding .input/.select raced and could leave the
    # LoRA prompt box hidden/empty while still applying defaults at generate time.
    # Pass live slider/prompt values + memory so existing settings survive add/remove.
    _lora_slider_inputs = [
        lora_selector, dynamic_loras_state,
        lora_weight_memory_state, lora_prev_selected_state,
        lora_prompt_memory_state, lora_prompt_display,
    ] + weight_sliders
    _lora_slider_outputs = (
        weight_sliders
        + [lora_prompt_display, selected_lora_details,
           lora_weight_memory_state, lora_prev_selected_state,
           lora_prompt_memory_state]
    )
    lora_selector.change(
        fn=update_weight_sliders,
        inputs=_lora_slider_inputs,
        outputs=_lora_slider_outputs,
        show_progress="hidden",
        trigger_mode="once",
    )

    # add_custom_lora is also imported from lora_registry.
    # Always session-only; optional persist / remove are separate explicit actions.
    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,
                 catalog_save_select, catalog_remove_select],
        show_progress="minimal",
    )
    catalog_save_select.change(
        fn=fill_catalog_save_form,
        inputs=[catalog_save_select, dynamic_loras_state],
        outputs=[catalog_save_title, catalog_save_weight, catalog_save_prompt,
                 catalog_save_triggers, catalog_save_notes],
        show_progress="hidden",
    )
    catalog_save_btn.click(
        fn=save_session_lora_to_catalog,
        inputs=[catalog_save_select, catalog_save_title, catalog_save_weight,
                catalog_save_prompt, dynamic_loras_state,
                catalog_save_triggers, catalog_save_notes,
                lora_selector],
        outputs=[catalog_save_status, lora_selector, catalog_save_select,
                 catalog_remove_select, dynamic_loras_state],
        show_progress="minimal",
    # After save, force weight/prompt UI to follow the remapped selection
    # (Custom: x → catalog title) so nothing stays bound to a removed title.
    ).then(
        fn=update_weight_sliders,
        inputs=_lora_slider_inputs,
        outputs=_lora_slider_outputs,
        show_progress="hidden",
    )
    catalog_remove_btn.click(
        fn=remove_lora,
        inputs=[catalog_remove_select, dynamic_loras_state, lora_selector],
        outputs=[catalog_remove_status, lora_selector, catalog_save_select,
                 catalog_remove_select, dynamic_loras_state],
        show_progress="minimal",
    )

    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],
        show_progress="hidden",
    )
    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],
        show_progress="hidden",
    )
    custom_prompt_selector.change(
        fn=update_custom_prompt_display,
        inputs=[custom_prompt_selector, custom_prompts_state],
        outputs=[custom_prompt_display],
        show_progress="hidden",
        trigger_mode="always_last",
    )

    canvas_mode.change(fn=on_canvas_mode_change, inputs=[canvas_mode],
                       outputs=[custom_width, custom_height], **_ui)
    canvas_fit_mode.change(fn=on_fit_mode_change, inputs=[canvas_fit_mode],
                           outputs=[pad_color], **_ui)
    batch_vary.change(fn=on_batch_vary_change, inputs=[batch_vary],
                      outputs=[sweep_min, sweep_max], **_ui)

    output_gallery.select(fn=on_gallery_select, inputs=[output_gallery],
                          outputs=[selected_output_state], show_progress="hidden")
    # Keep download buttons pointed at the selected gallery item (or latest)
    # plus the last-run prompt text.
    selected_output_state.change(
        fn=resolve_download_bundle,
        inputs=[selected_output_state, output_gallery, last_prompt_state, project_name],
        outputs=[download_png_btn, download_webp_btn, download_prompt_btn],
        show_progress="hidden",
    )
    output_gallery.change(
        fn=resolve_download_bundle,
        inputs=[selected_output_state, output_gallery, last_prompt_state, project_name],
        outputs=[download_png_btn, download_webp_btn, download_prompt_btn],
        show_progress="hidden",
    )
    # Track which gallery images still need downloading.
    output_gallery.change(
        fn=sync_download_tracking,
        inputs=[output_gallery, pending_download_state, downloaded_images_state],
        outputs=[pending_download_state, downloaded_images_state, download_status],
        show_progress="hidden",
    )
    project_name.change(
        fn=resolve_download_bundle,
        inputs=[selected_output_state, output_gallery, last_prompt_state, project_name],
        outputs=[download_png_btn, download_webp_btn, download_prompt_btn],
        show_progress="hidden",
    )
    # Open PNG/WebP in a new tab (in addition to the browser download) and mark
    # the current gallery image as downloaded for the pending-status tracker.
    _OPEN_DL_TAB_JS = """
(btnId) => {
  const openHref = (href) => {
    if (!href || href === '#' || href.endsWith('/')) return false;
    window.open(href, '_blank', 'noopener,noreferrer');
    return true;
  };
  const tryOpen = () => {
    const root = document.getElementById(btnId);
    if (!root) return false;
    const anchors = root.querySelectorAll('a.download-link, a[href], a[download]');
    for (const a of anchors) {
      const href = a.href || a.getAttribute('href') || '';
      if (openHref(href)) return true;
    }
    // Gradio sometimes nests the file link one tick later after value bind.
    return false;
  };
  if (tryOpen()) return;
  // Retry briefly — DownloadButton href can lag the click on first bind.
  let n = 0;
  const t = setInterval(() => {
    n += 1;
    if (tryOpen() || n >= 8) clearInterval(t);
  }, 50);
}
"""
    download_png_btn.click(
        fn=None,
        # Concatenate (not f-string) so braces inside _OPEN_DL_TAB_JS stay literal JS.
        js="() => { (" + _OPEN_DL_TAB_JS + ")('download_png_btn'); }",
    ).then(
        fn=mark_current_output_downloaded,
        inputs=[selected_output_state, output_gallery,
                pending_download_state, downloaded_images_state],
        outputs=[pending_download_state, downloaded_images_state, download_status],
        show_progress="hidden",
    )
    download_webp_btn.click(
        fn=None,
        js="() => { (" + _OPEN_DL_TAB_JS + ")('download_webp_btn'); }",
    ).then(
        fn=mark_current_output_downloaded,
        inputs=[selected_output_state, output_gallery,
                pending_download_state, downloaded_images_state],
        outputs=[pending_download_state, downloaded_images_state, download_status],
        show_progress="hidden",
    )
    # Gallery built-in ↓ icon → JS signal → mark downloaded.
    gallery_dl_signal.change(
        fn=mark_from_gallery_signal,
        inputs=[gallery_dl_signal, output_gallery,
                pending_download_state, downloaded_images_state],
        outputs=[pending_download_state, downloaded_images_state,
                 download_status, gallery_dl_signal],
        show_progress="hidden",
    )

    # Reset any stale gallery selection before a new run starts, so Send→Base
    # / Send→Ref after this run can't accidentally reuse a path from the
    # previous run's gallery contents. Also warn if undownloaded images exist.
    _infer_inputs = [
        base_image, ref1, ref2, ref3, 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, project_name,
        dynamic_loras_state,
    ] + weight_sliders
    _infer_outputs = [
        output_gallery, used_seed,
        download_png_btn, download_webp_btn, download_prompt_btn,
        last_prompt_state,
    ]
    for _run_btn in (run_button, run_button_top):
        _run_btn.click(
            fn=warn_undownloaded_before_generate,
            inputs=[pending_download_state, output_gallery],
            outputs=[selected_output_state],
            show_progress="hidden",
        )

    run_event = run_button.click(
        fn=infer,
        inputs=_infer_inputs,
        # Download bundle on every generate yield — required because
        # gallery.change alone misses the first virgin-session result.
        outputs=_infer_outputs,
    )
    run_event_top = run_button_top.click(
        fn=infer,
        inputs=_infer_inputs,
        outputs=_infer_outputs,
    )

    # ── Editor tab wiring ────────────────────────────────────────────────────
    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_ref_outputs = [
        ref1, ref2, ref3, reference_info,
        ref2_up_btn, ref2_down_btn, ref3_up_btn,
    ]
    send_to_ref_btn.click(
        fn=send_editor_to_reference,
        inputs=[editor, ref1, ref2, ref3],
        outputs=_send_ref_outputs,
    ).then(fn=lambda: gr.Tabs(selected="tab_generate"), outputs=[main_tabs])

    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, ref1, ref2, ref3],
        outputs=_send_ref_outputs,
    )

    # Extend-canvas wiring
    # The schematic previews the *current editor composite* so users see live
    # feedback as they nudge the percentages / fill colour.
    def _editor_source_for_preview(editor_value):
        if not editor_value or editor_value.get("composite") is None:
            return None
        comp = editor_value["composite"]
        if isinstance(comp, np.ndarray):
            from PIL import Image as _Image
            comp = _Image.fromarray(comp)
        return comp

    def _update_extend_preview(editor_value, up, down, left, right, fill):
        return render_extend_schematic(
            _editor_source_for_preview(editor_value), up, down, left, right, fill,
        )

    _extend_preview_inputs = [editor, ext_up, ext_down, ext_left, ext_right, ext_fill]
    _extend_preview_outputs = [ext_schematic, ext_info]
    for _c in (ext_up, ext_down, ext_left, ext_right, ext_fill):
        _c.change(fn=_update_extend_preview,
                  inputs=_extend_preview_inputs, outputs=_extend_preview_outputs)
    # Refresh the schematic when a NEW image lands in the editor — not on every
    # `change` event. `editor.change` fires very frequently on iOS Safari/Chrome
    # (once per stroke/layer/crop-preview) and the round-trips OOM'd the tab
    # even for small uploads. `.upload` fires only when a new image comes in
    # via the upload/clipboard sources, which is the case the preview actually
    # cares about (image dimensions changed → schematic scale needs redrawing).
    editor.upload(fn=_update_extend_preview,
                  inputs=_extend_preview_inputs, outputs=_extend_preview_outputs)
    # HEIC uploads bypass the editor's own upload event because they come from
    # the separate File component, so wire that path in explicitly too.
    heic_uploader.upload(fn=_update_extend_preview,
                         inputs=_extend_preview_inputs, outputs=_extend_preview_outputs)

    # Run: extend, then hand the new PIL back to the editor. `render_extend_
    # schematic` re-runs via editor.change once the new image lands, so no
    # extra .then() is needed for the preview.
    extend_btn.click(
        fn=extend_editor_canvas,
        inputs=[editor, ext_up, ext_down, ext_left, ext_right, ext_fill],
        outputs=[editor],
    )

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

    bulk_stop_btn.click(fn=lambda: gr.Info("Stop requested — finishing current image."),
                        cancels=[bulk_event, run_event, run_event_top])

    # ── Depth / Pose tab wiring ──────────────────────────────────────────────

    ctrl_source.change(
        fn=lambda img: img, inputs=[ctrl_source], outputs=[pose_source_state],
    )

    detect_depth_btn.click(
        fn=generate_depthmap, inputs=[ctrl_source], outputs=[depth_output],
    )

    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(
        fn=_on_detect_pose, inputs=[ctrl_source],
        outputs=[pose_keypoints_state, active_person_dd, active_joint_dd,
                 pose_overlay, pose_clean],
    )

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

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

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

    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_depth_ref_btn.click(
        fn=push_pil_to_reference, inputs=[depth_output, ref1, ref2, ref3],
        outputs=_send_ref_outputs,
    ).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, ref1, ref2, ref3],
        outputs=_send_ref_outputs,
    ).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 on launch(), not Blocks()
    demo.queue().launch(css=css, theme=orange_red_theme,
                        mcp_server=True, ssr_mode=False, show_error=True)