app.py

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

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,
    ENABLE_LOGGING,
)
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,
)
from logging_utils import log_inference
from image_utils import (
    fix_orientation,
    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,
    extend_editor_canvas,
    render_extend_schematic,
)
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,
)

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


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

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=""):
    """Fire-and-forget logger. ENABLE_LOGGING is read once at import time in
    config.py — there's no per-call env lookup. Used by single, batch and
    bulk inference paths."""
    if not ENABLE_LOGGING:
        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, 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


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

    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.

    Logging: each image is logged individually via _spawn_log on success or
    failure — same code path as single/batch, so bulk shows up in the dataset
    with the same schema, just with extra bulk_index / bulk_total fields
    embedded in the PNG metadata."""
    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,
            )

            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()
        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}"])
            _spawn_log([], None, prompt, 0, steps, guidance_scale,
                       0, 0, duration, 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"⚠️ Image {i+1} failed: {e} | {succeeded} ok, {failed} failed", gr.update()
            continue

    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


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

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)

    with gr.Column(elem_id="col-container"):
        _logging_badge = "🟢 On" if ENABLE_LOGGING 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)

                    with gr.Column(scale=1):
                        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).*"
                        )

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

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

    # update_weight_sliders is the one imported from lora_registry now.
    lora_selector.change(
        fn=update_weight_sliders,
        inputs=[lora_selector, dynamic_loras_state],
        outputs=weight_sliders + [lora_prompt_display],
    )

    # add_custom_lora is also imported from lora_registry.
    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_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])

    output_gallery.select(fn=on_gallery_select, inputs=[output_gallery],
                          outputs=[selected_output_state])
    # 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.
    run_button.click(fn=lambda: None, outputs=[selected_output_state])

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

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

    # ── 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, 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 on launch(), not Blocks()
    demo.queue().launch(css=css, theme=orange_red_theme,
                        mcp_server=True, ssr_mode=False, show_error=True)