"""MiniMax-H3 `ref2va`, split deployment — the denoising half.

This Space holds the `transformer_ref` partition and the two autoencoders, unquantized bfloat16. Text encoding runs in
[`qwen3vl-conditioner`](https://huggingface.co/spaces/multimodalart/qwen3vl-conditioner), which this one calls over the
gradio API for every request; `reference_encoder` stays here, next to the autoencoders it runs.
"""

from __future__ import annotations

import json
import os
import random
import re
import subprocess
import tempfile
import time
import traceback
from functools import cache

# Before anything that could initialize CUDA: `import spaces` patches `torch.cuda` so the 72 GiB load can happen at
# startup rather than on GPU time.
import spaces
import gradio as gr

MODEL_REPO = os.environ.get("H3_MODEL_REPO", "MiniMaxAI/MiniMax-H3")
CONDITIONER_SPACE = os.environ.get("H3_CONDITIONER", "multimodalart/qwen3vl-conditioner")
# `lazy` moves all 72.16 GiB onto the card on the first GPU call and leaves it there; `offload` hands placement to
# `ComponentsManager.enable_auto_cpu_offload`. Startup placement is not an option here — see `load_models`.
PLACEMENT = os.environ.get("H3_PLACEMENT", "lazy").lower()
# cuDNN's fused attention is 10-20% faster than the SDPA default on this pool and needs nothing installed.
# flash-attention 3 is sm90-only and this card is sm120 (the `zero-a10g` flavour name is legacy).
ATTENTION = os.environ.get("H3_ATTENTION", "_native_cudnn").lower()
GPU_SIZE = os.environ.get("H3_GPU_SIZE", "xlarge")
# Bounds on what `get_duration` may reserve. The pool reserves whatever number it is given, so a flat ceiling for every
# request is what makes an account hit "too many ZeroGPU credits allocated to running tasks".
MIN_GPU_DURATION = int(os.environ.get("H3_GPU_DURATION_MIN", "120"))
MAX_GPU_DURATION = int(os.environ.get("H3_GPU_DURATION_MAX", "1500"))
# Ceiling on the packed sequence. Above it the card dies inside the rotary embeddings with
# `NVML_SUCCESS == r INTERNAL ASSERT FAILED` - not a bug in the code, just out of memory.
# 74k rows go through; 165k kill the worker.
MAX_SEQUENCE = int(os.environ.get("H3_MAX_SEQUENCE", "90000"))
# How many text rows the pre-flight estimate allows for, before the conditioner returns the exact count.
TEXT_TOKEN_ALLOWANCE = int(os.environ.get("H3_TEXT_TOKEN_ALLOWANCE", "13000"))

# Must stay identical to the conditioner's table: the *label* goes over the wire, so a canvas that half does not know
# is rejected there and surfaces as a failure here.
CANVASES = {
    # 16:9
    "960x544 · 16:9 fast": (544, 960),
    "1024x576 · 16:9 fast": (576, 1024),
    "1152x640 · 16:9": (640, 1152),
    "1280x704 · 16:9": (704, 1280),
    "1344x768 · 16:9 full": (768, 1344),
    # 9:16
    "544x960 · 9:16 fast": (960, 544),
    "640x1152 · 9:16": (1152, 640),
    "768x1344 · 9:16 full": (1344, 768),
    # 1:1
    "544x544 · 1:1 fast": (544, 544),
    "768x768 · 1:1 full": (768, 768),
    # 4:3 / 3:4
    "768x576 · 4:3 fast": (576, 768),
    "1024x768 · 4:3 full": (768, 1024),
    "576x768 · 3:4 fast": (768, 576),
    "768x1024 · 3:4 full": (1024, 768),
    # 21:9
    "1152x512 · 21:9 fast": (512, 1152),
    "1536x672 · 21:9 full": (672, 1536),
}
DEFAULT_CANVAS = "960x544 · 16:9 fast"
FPS, FRAMES_PER_CHUNK, LATENTS_PER_CHUNK = 24, 17, 5
# It is the *snapped* frame count the ceiling has to hold for: 15 s is 360 frames, which rounds up to 362, i.e.
# 15.083 s, and is refused. 14 is the last whole second that survives the snap.
MAX_UI_DURATION = 14
MIN_DURATION = 2
# A reference video shorter than 2 s gives the model almost no motion to read.
MIN_REFERENCE_VIDEO, MAX_REFERENCE_VIDEO = 2.0, 15.0
# `MINIMAX_H3_MAX_REFERENCE_IMAGES`. The slots are built up front and revealed one at a time, because a demo asking
# for two subjects should not open with nine boxes.
MAX_IMAGE_SLOTS, OPEN_IMAGE_SLOTS = 9, 2

# How many LoRA slots the UI offers, and the range each strength slider covers. Everything else -
# the UI loop, the settings keys, the preset filler, `generate`'s `*lora_fields` tail - is built from
# this number, so it is the only place to change it.
LORA_SLOTS = 5
LORA_MIN_SCALE, LORA_MAX_SCALE = -2.0, 2.0
# What an empty slot starts at. Most H3 adapters on CivitAI are written up for 0.5, and stacking two or
# three of them at 1.0 is what turns a clip plastic. A Turbo preset fills its slot with its own strength
# instead of this one.
DEFAULT_LORA_SCALE = 0.5

# Pre-wired Turbo LoRAs from `larryvrh/MiniMax-H3-Turbo-Lora`: a few-step distillation that renders joint video +
# soundtrack in 4–8 steps instead of the usual ~20. Each entry is `(repo reference, recommended steps, blurb)`. The
# reference is the `owner/repo/filename.safetensors` form `resolve_lora` accepts, so it downloads on first use and is
# cached by `huggingface_hub` thereafter — nothing is bundled in this Space.
#
# The fourth element is the strength the slot is filled at. Larryvrh documents 1.0 for every build, and that is
# what v4 gets; the older v1 line is the one people report over-sharpening on at strength 1.0, so it is filled at
# 0.7 instead. Both are starting points - the slider is right there.
LORA_PRESETS = {
    "Turbo v4 step600 EMA · 8 steps (recommended)": (
        "larryvrh/MiniMax-H3-Turbo-Lora/minimax_h3_turbo_v4_step600_ema.safetensors",
        8,
        "The author's own recommendation and the strongest build released: much better static and small-motion "
        "shots, markedly better micro-detail in faces, fingers and fine texture, and the plastic over-sharpened "
        "look of the older v1 line is gone. Its one weak spot is 4 steps with large fast motion, where it can "
        "trail - 6 to 8 steps removes that and is where it looks its best.",
        1.0,
    ),
    "Turbo v4 step600 non-EMA · 6 steps": (
        "larryvrh/MiniMax-H3-Turbo-Lora/minimax_h3_turbo_v4_step600.safetensors",
        6,
        "The same training run without the EMA averaging. The author recommends the EMA build, but a number of "
        "users report cleaner results from this one - worth a try if EMA output looks soft.",
        1.0,
    ),
    "Turbo v1 ckpt850 · 4 steps (fast motion only)": (
        "larryvrh/MiniMax-H3-Turbo-Lora/minimax_h3_turbo_4step_ema_ckpt850.safetensors",
        4,
        "Superseded by v4 in every other respect, and kept for one case the author names: at 4 steps with large "
        "fast motion, v4 trails and this older build does not. Filled at 0.7 because the v1 line over-sharpens "
        "at 1.0. Anything that is not fast motion at 4 steps belongs on v4.",
        0.7,
    ),
}

# Every Turbo build tunes the *video* trajectory. The soundtrack has its own (flow shift 12 for video against 3 for
# audio), which Larryvrh's ComfyUI Turbo *sampler* handles and this diffusers Space does not have - so at 4 steps the
# audio can come out distorted even when the picture is fine. Raise the steps if it does.

# The lowest step count the model's own schedulers accept; the Turbo LoRAs are tuned for 4.
MIN_STEPS = 4

# Seconds of GPU one request needs, from the packed sequence it is about to denoise: linear in the rows for the
# matmuls, quadratic for the attention, against the AoTI block package this Space runs.
STEP_LINEAR, STEP_QUADRATIC, SAFETY = 1.1745e-4, 3.8396e-9, 1.3
# The lazy 72.16 GiB `PIPE.to("cuda")` a cold worker pays inside its first GPU call; every request carries it, because
# nothing here knows whether the worker it lands on is cold.
PLACEMENT_ALLOWANCE = int(os.environ.get("H3_PLACEMENT_ALLOWANCE", "90"))
AUDIO_LATENTS_PER_SECOND, AUDIO_CHANNELS = 40, 2
REFERENCE_IMAGE_SHORT_EDGE, CANVAS_MULTIPLE = 2048, 32
DECODE_BASE, DECODE_PER_DEFAULT_CANVAS, DEFAULT_CANVAS_PIXELS = 15, 25, 960 * 544 * 124
# Reading one adapter off local disk and injecting it across the 33B transformer's linear layers.
LORA_ALLOWANCE = 12


def snap_frames(seconds: float) -> int:
    """The frame count MiniMax-H3's video VAE can decode: the next `17 * n + 5` at 24 fps."""
    frames = max(1, round(float(seconds) * FPS))
    while frames % FRAMES_PER_CHUNK != LATENTS_PER_CHUNK:
        frames += 1
    return frames


def lower_duration_floor(seconds: float = MIN_DURATION) -> None:
    """Let the pipeline generate below its 5 s floor. 56 frames (2.33 s) is fine on the released checkpoint."""
    from diffusers.modular_pipelines.minimax_h3.modular_pipeline import MiniMaxH3ModularPipeline

    MiniMaxH3ModularPipeline.min_duration = property(lambda self: float(seconds))


def video_latent_frames(num_frames: int) -> int:
    """`17 * n + 5` frames become `5 * n + 2` video latents."""
    return 5 * ((num_frames - LATENTS_PER_CHUNK) // FRAMES_PER_CHUNK) + 2


def target_rows(height: int, width: int, num_frames: int) -> int:
    """The generated rows of the packed sequence: video patched `(1, 2, 2)`, plus two audio rows per latent."""
    video = video_latent_frames(num_frames) * (height // CANVAS_MULTIPLE) * (width // CANVAS_MULTIPLE)
    return video + round(num_frames / FPS * AUDIO_LATENTS_PER_SECOND) * AUDIO_CHANNELS


def reference_rows(references: list[tuple[str, str]], num_frames: int) -> int:
    """The rows the reference blocks add, from metadata alone — no decode.

    An image is resized to a 2048 pixel short edge and encoded as a single frame; a video is put on the canvas *its
    own* aspect ratio resolves to, truncated to the generated frame count and snapped **down** to a `17 * n + 5` the
    VAE encodes without padding; a soundtrack contributes two rows per 1/40 s.
    """
    from PIL import Image

    from diffusers.modular_pipelines.minimax_h3.modular_pipeline import resolve_canvas_size

    rows = 0
    for kind, path in references:
        if kind == "image":
            width, height = Image.open(path).size
            scale = REFERENCE_IMAGE_SHORT_EDGE / min(width, height)
            resolved = [
                max(CANVAS_MULTIPLE, round(edge * scale / CANVAS_MULTIPLE) * CANVAS_MULTIPLE)
                for edge in (height, width)
            ]
            rows += (resolved[0] // CANVAS_MULTIPLE) * (resolved[1] // CANVAS_MULTIPLE)
            continue

        video_seconds, audio_seconds = probe(path)
        if kind == "video" and video_seconds is not None:
            import av

            with av.open(path) as container:
                stream = container.streams.video[0]
                source_height, source_width = stream.height, stream.width
            canvas_height, canvas_width = resolve_canvas_size(source_width, source_height, CANVAS_MULTIPLE)
            frames = min(round(video_seconds * FPS), num_frames)
            snapped = max(1, (frames - LATENTS_PER_CHUNK) // FRAMES_PER_CHUNK) * FRAMES_PER_CHUNK + LATENTS_PER_CHUNK
            rows += (
                video_latent_frames(snapped)
                * (canvas_height // CANVAS_MULTIPLE)
                * (canvas_width // CANVAS_MULTIPLE)
            )
        if audio_seconds is not None:
            seconds = min(audio_seconds, num_frames / FPS)
            rows += round(seconds * AUDIO_LATENTS_PER_SECOND) * AUDIO_CHANNELS
    return rows


def get_duration(
    prompt_embeds, text_token_tags, references, height, width, num_frames, steps, seed, loras=(), **_
):
    """Seconds of GPU to reserve for one request. Takes the arguments of the `@spaces.GPU` function it decorates, and
    tolerates the `gr.Progress` `spaces` injects."""
    sequence = int(text_token_tags.shape[0]) + reference_rows(references, num_frames) + target_rows(
        height, width, num_frames
    )
    denoise = int(steps) * (STEP_LINEAR * sequence + STEP_QUADRATIC * sequence**2) * SAFETY
    # The two reference encoders ahead of the loop, and the two decoders plus the mux after it. Both scale with what
    # they are handed rather than with the step count.
    encode = 5 + reference_rows(references, num_frames) * 1e-3
    decode = DECODE_BASE + DECODE_PER_DEFAULT_CANVAS * (height * width * num_frames) / DEFAULT_CANVAS_PIXELS
    total = PLACEMENT_ALLOWANCE + encode + denoise + decode + 10 + LORA_ALLOWANCE * len(loras or ())
    duration = max(MIN_GPU_DURATION, min(MAX_GPU_DURATION, int(total)))
    print(f"[ref2va] S={sequence} -> reserving {duration}s ({denoise:.0f}s of denoise at {steps} steps)", flush=True)
    return duration


def budget(text_tokens, references, height, width, num_frames, steps, loras=()):
    """`(rows, GPU seconds)` for one request, by the same formula as `get_duration`."""
    sequence = int(text_tokens) + reference_rows(references, num_frames) + target_rows(height, width, num_frames)
    per_step = (STEP_LINEAR * sequence + STEP_QUADRATIC * sequence**2) * SAFETY
    encode = 5 + reference_rows(references, num_frames) * 1e-3
    decode = DECODE_BASE + DECODE_PER_DEFAULT_CANVAS * (height * width * num_frames) / DEFAULT_CANVAS_PIXELS
    overhead = PLACEMENT_ALLOWANCE + encode + decode + 10 + LORA_ALLOWANCE * len(loras or ())
    return sequence, overhead + int(steps) * per_step, per_step, overhead


def fits(text_tokens, references, height, width, num_frames, steps, loras=()):
    """Stops a request the card or the reservation cannot take, before any GPU time is spent."""
    sequence, total, per_step, overhead = budget(
        text_tokens, references, height, width, num_frames, steps, loras
    )
    if sequence <= MAX_SEQUENCE and total <= MAX_GPU_DURATION:
        return

    seconds = num_frames / FPS
    if sequence > MAX_SEQUENCE:
        raise gr.Error(
            f"This request is too large for the card: {sequence} rows against a ceiling of {MAX_SEQUENCE} "
            f"({width}x{height}, {seconds:.1f} s, {len(references)} references). "
            "Lower the duration, pick a smaller canvas, or remove a reference."
        )

    room = int((MAX_GPU_DURATION - overhead) / per_step)
    advice = (
        f"Lower Steps to {room}." if room >= MIN_STEPS
        else "Lower the duration or pick a smaller canvas."
    )
    raise gr.Error(
        f"This request wants ~{int(total)} s of GPU, and the ceiling is {MAX_GPU_DURATION} s "
        f"({width}x{height}, {seconds:.1f} s, {int(steps)} steps). {advice}"
    )


PIPE = None
MANAGER = None
LOAD_ERROR: str | None = None


def load_models() -> str | None:
    """Load the denoising half at startup, but *not* onto the card.

    `MiniMaxH3Ref2VAGeneratorBlocks` declares `transformer_ref`, `vae`, `audio_vae`, the two schedulers and
    `video_processor`, so `load_components` fetches exactly those subfolders — `text_encoder/` and the `transformer/`
    partition are never touched. Both autoencoders carry `_keep_in_fp32_modules` over every module and stay float32: a
    bfloat16 audio VAE decodes the soundtrack roughly 20 dB too quiet.

    Nothing moves onto the card here, for storage rather than memory: `spaces`' startup `torch.pack()` writes every
    startup-resident CUDA tensor to a second copy on disk, and 77.3 GB of weights plus its pack busts the 150 GB quota
    (`OSError: [Errno 28] No space left on device` out of `os.posix_fallocate`, mid-pack).
    """
    global PIPE, MANAGER, LOAD_ERROR

    if PIPE is not None or LOAD_ERROR is not None:
        return LOAD_ERROR

    started = time.time()
    try:
        import torch
        from diffusers import ComponentsManager

        from h3_split_blocks import MiniMaxH3Ref2VAGeneratorBlocks

        lower_duration_floor()
        manager = ComponentsManager()
        blocks = MiniMaxH3Ref2VAGeneratorBlocks()
        print(f"[ref2va] loading {[c.name for c in blocks.expected_components]} from {MODEL_REPO} ...", flush=True)
        pipe = blocks.init_pipeline(MODEL_REPO, components_manager=manager, collection="h3")
        pipe.load_components(dtype=torch.bfloat16)

        # Both VAEs first, and explicitly. `set_attention_backend` also sets the registry's *global* backend, which
        # every processor that was not stamped falls through to, and the float32 audio VAE has no cuDNN kernel:
        # `RuntimeError: No available kernel. Aborting execution.` in its causal encoder attention, which only a
        # reference soundtrack ever reaches.
        pipe.vae.set_attention_backend("native")
        pipe.audio_vae.set_attention_backend("native")
        pipe.transformer_ref.set_attention_backend(ATTENTION)

        # Still startup, still free: an AoTI package carries no weights and opens its archive lazily inside the GPU
        # worker. Off unless `H3_AOTI=1`. It is the *same* package the `transformer/` partition runs — the two configs
        # are identical field for field and the compiled code carries no weights of either.
        import h3_aoti

        h3_aoti.maybe_load(pipe.transformer_ref)

        if PLACEMENT == "offload":
            manager.enable_auto_cpu_offload(device="cuda")
            _arm_decode_hooks(pipe)

        PIPE, MANAGER = pipe, manager
        print(f"[ref2va] ready in {time.time() - started:.0f}s", flush=True)
    except Exception as error:
        traceback.print_exc()
        LOAD_ERROR = (
            f"**Loading `{MODEL_REPO}` failed** after {time.time() - started:.0f}s: "
            f"`{type(error).__name__}: {error}`"
        )
    return LOAD_ERROR


def _arm_decode_hooks(pipe):
    """Make the offload hooks fire for the two VAEs.

    `enable_auto_cpu_offload` wraps `forward`, and the reference-encoder and decode blocks call `vae.encode/decode(...)`
    directly, so the hook never runs and the VAE is still on the host when the latents arrive on the card.
    """
    for name in ("vae", "audio_vae"):
        module = getattr(pipe, name)
        for method in ("encode", "decode"):
            inner = getattr(module, method)

            def armed(*args, _module=module, _inner=inner, **kwargs):
                hook = getattr(_module, "_hf_hook", None)
                if hook is not None:
                    hook.pre_forward(_module)
                return _inner(*args, **kwargs)

            setattr(module, method, armed)


# ----------------------------------------------------------------------------------------------------------------
# LoRA
# ----------------------------------------------------------------------------------------------------------------
# There is no `MiniMaxH3LoraLoaderMixin` in the diffusers integration, so adapters are attached at the *model* level,
# through the `PeftAdapterMixin` the transformer carries. That is the whole API this needs: `load_lora_adapter` for
# each file and one `set_adapters` call to give them their strengths. Here the model is `transformer_ref`, so the
# adapters have to be trained against the `transformer_ref/` partition — a `transformer/` adapter is a different
# partition and will not match.


def _hub_url_parts(url: str) -> tuple[str, str]:
    """Split a huggingface.co `blob`/`resolve` URL into its repo id and the file path inside it."""
    from urllib.parse import unquote, urlparse

    parts = unquote(urlparse(url).path).strip("/").split("/")
    if len(parts) < 5 or parts[2] not in ("resolve", "blob"):
        raise gr.Error(f"That address is not a recognisable Hugging Face file URL: `{url}`")
    return "/".join(parts[:2]), "/".join(parts[4:])


CIVITAI_HOSTS = ("civitai.com", "civitai.red", "civitai.green", "civitai.work")


def _download_direct_lora(url: str) -> str:
    """Fetch a `.safetensors` from a plain URL - CivitAI in particular - and return the local path.

    The download happens on the Space's own machine, not in the visitor's browser, so a CivitAI
    session in a browser tab has nothing to do with it: a gated model answers a server with an HTML
    login page instead of weights. `CIVITAI_TOKEN` (Settings -> Variables and secrets) is appended
    automatically when it is set, and the header of whatever comes back is checked so a login page
    fails with a sentence that says what to do rather than a parse error deep inside safetensors.
    """
    import hashlib
    from urllib.parse import urlparse, unquote

    import requests

    host = (urlparse(url).hostname or "").lower()
    request_url = url
    token = os.environ.get("CIVITAI_TOKEN", "").strip()
    if token and any(host.endswith(known) for known in CIVITAI_HOSTS) and "token=" not in url:
        request_url = url + ("&" if "?" in url else "?") + f"token={token}"

    cache_dir = os.path.join(tempfile.gettempdir(), "url-loras")
    os.makedirs(cache_dir, exist_ok=True)
    cached = os.path.join(cache_dir, hashlib.sha256(url.encode()).hexdigest()[:16] + ".safetensors")
    if os.path.exists(cached) and os.path.getsize(cached) > 1_000_000:
        return cached

    try:
        response = requests.get(request_url, stream=True, timeout=120,
                                headers={"User-Agent": "Mozilla/5.0"})
        response.raise_for_status()
    except Exception as error:
        raise gr.Error(f"Could not download `{url}`: {error}")

    content_type = (response.headers.get("content-type") or "").lower()
    if "text/html" in content_type:
        raise gr.Error(
            "That link answered with a web page instead of a file. The model is gated, so the "
            "Space needs its own key: add CIVITAI_TOKEN under Settings -> Variables and secrets."
        )

    disposition = response.headers.get("content-disposition", "")
    name = unquote(re.findall(r'filename\*?=(?:UTF-8\'\'|")?([^";]+)', disposition)[0]) \
        if "filename" in disposition else os.path.basename(urlparse(url).path)
    if name and not name.lower().endswith(".safetensors") and "." in name:
        print(f"[lora] {name} is not a .safetensors; trying it anyway")

    written = 0
    with open(cached, "wb") as handle:
        for chunk in response.iter_content(chunk_size=1 << 20):
            if chunk:
                handle.write(chunk)
                written += len(chunk)

    if written < 1_000_000:
        os.remove(cached)
        raise gr.Error(
            "That link returned only a few kilobytes - almost always a login or error page rather "
            "than weights. Check the link, or add CIVITAI_TOKEN to the Space."
        )

    # safetensors starts with an 8-byte little-endian header length followed by that much JSON.
    with open(cached, "rb") as handle:
        header_len = int.from_bytes(handle.read(8), "little")
        if not (0 < header_len < 100_000_000):
            os.remove(cached)
            raise gr.Error("The downloaded file is not a `.safetensors` (bad header).")
        try:
            json.loads(handle.read(header_len).decode("utf-8"))
        except Exception:
            os.remove(cached)
            raise gr.Error("The downloaded file is not a `.safetensors` (unreadable header).")

    print(f"[lora] downloaded {written / 1e6:.0f} MB from {host} -> {os.path.basename(cached)}")
    return cached


def resolve_lora(reference: str) -> str:
    """Turn what the user typed into a local `.safetensors` path.

    Accepts a local path, a huggingface.co file URL, `owner/repo/path/to/file.safetensors`, or a bare `owner/repo`
    whose single `.safetensors` is then picked for them. Runs outside the GPU call, so the download costs no GPU time.
    """
    from huggingface_hub import hf_hub_download, list_repo_files

    reference = (reference or "").strip()
    if not reference:
        return ""
    if os.path.exists(reference):
        return reference
    if reference.startswith(("http://", "https://")):
        from urllib.parse import urlparse

        if (urlparse(reference).hostname or "").lower().endswith("huggingface.co"):
            repo_id, filename = _hub_url_parts(reference)
            return hf_hub_download(repo_id, filename)
        # Anything else - CivitAI and any other direct link - is fetched as a plain file.
        return _download_direct_lora(reference)

    parts = [part for part in reference.split("/") if part]
    if len(parts) > 2 and parts[-1].endswith(".safetensors"):
        return hf_hub_download("/".join(parts[:2]), "/".join(parts[2:]))
    if len(parts) != 2:
        raise gr.Error(
            f"`{reference}` is not an existing file, an `owner/repo`, or a Hugging Face URL."
        )

    candidates = [name for name in list_repo_files(reference) if name.endswith(".safetensors")]
    if not candidates:
        raise gr.Error(f"`{reference}` holds no `.safetensors` file.")
    if len(candidates) > 1:
        preferred = [name for name in candidates if "lora" in name.lower()]
        if len(preferred) != 1:
            listed = ", ".join(f"`{name}`" for name in sorted(candidates)[:8])
            raise gr.Error(f"`{reference}` holds several files. Write `{reference}/name.safetensors`. Available: {listed}")
        candidates = preferred
    return hf_hub_download(reference, candidates[0])


def _lora_prefix(state_dict) -> str | None:
    """The prefix `load_lora_adapter` has to strip before the keys match the transformer's own module names."""
    key = next(iter(state_dict))
    for prefix in ("model.diffusion_model", "diffusion_model", "transformer_ref", "transformer"):
        if key.startswith(f"{prefix}."):
            return prefix
    return None


# ------------------------------------------------------------------------------------------------------------------
# CivitAI / kohya LoRA conversion
# ------------------------------------------------------------------------------------------------------------------
# Adapters trained with kohya-style trainers - which is most of what CivitAI carries - differ from diffusers in more
# ways than the ComfyUI Turbo LoRA does, and each of the three below silently ruins the result rather than raising:
#
#   * names are flat and underscored (`lora_unet_blocks_0_attn_qkv_proj`) rather than dotted,
#   * the fused QKV is interleaved *per attention head* (q,k,v for head 0, then head 1, ...), not three plain thirds,
#     so splitting it with `chunk(3)` hands q's rows to k and k's to v,
#   * the gated MLP's `fc1` keeps its two halves in the opposite order to diffusers' `ff.net.0.proj`,
#   * `alpha` sets the scale, and ignoring it makes the adapter arrive at the wrong strength.
#
# Ported from the standalone converter, so a CivitAI file can be pasted straight into a slot.

NUM_HEADS, HEAD_DIM = 56, 128
INNER_DIM = NUM_HEADS * HEAD_DIM  # 7168

_DOT_PREFIXES = (
    "base_model.model.", "base_model.", "model.diffusion_model.", "diffusion_model.", "transformer.", "net.",
)
_FLAT_PREFIXES = ("lora_unet_", "lora_transformer_", "lora_te_", "lora_")
_LORA_SUFFIXES = (
    (".lora_down.weight", "down"), (".lora_up.weight", "up"),
    (".lora_A.weight", "down"), (".lora_B.weight", "up"),
    (".lora_A", "down"), (".lora_B", "up"),
    (".alpha", "alpha"), (".lora_alpha", "alpha"),
)
_DIFFUSERS_MARKERS = (".to_q", ".to_k", ".to_v", ".to_out.0", "ff.net.0.proj", "transformer_blocks.")
_KOHYA_NAME = re.compile(r"^(token_refiner_)?blocks_(\d+)_(attn_qkv_proj|attn_out_proj|mlp_fc1|mlp_fc2)$")


def _strip_lora_prefixes(name: str) -> str:
    changed = True
    while changed:
        changed = False
        for prefix in _DOT_PREFIXES + _FLAT_PREFIXES:
            if name.startswith(prefix):
                name, changed = name[len(prefix):], True
    return name


def _dotted_module(name: str):
    """`blocks_12_attn_qkv_proj` -> `blocks.12.attn.qkv_proj`. Already-dotted names pass through."""
    if "blocks." in name:
        return name
    match = _KOHYA_NAME.match(name)
    if not match:
        return None
    refiner, index, leaf = match.groups()
    head = "token_refiner.blocks." if refiner else "blocks."
    return f"{head}{index}.{leaf.replace('_', '.', 1)}"


def _rename_module(module: str) -> str:
    if module.startswith("token_refiner.blocks."):
        module = module.replace("token_refiner.blocks.", "token_refiner.refiner_blocks.", 1)
    elif module.startswith("blocks."):
        module = module.replace("blocks.", "transformer_blocks.", 1)
    return module.replace(".attn.out_proj", ".attn.to_out.0").replace(".mlp.fc2", ".ff.net.2")


def _split_qkv_by_head(tensor):
    """Undo the per-head interleave of a fused QKV `lora_B` of shape `[3 * INNER_DIM, rank]`.

    Row order is head 0's q, k and v, then head 1's, and so on, so taking three contiguous thirds is wrong; the rows
    have to be gathered a head at a time.
    """
    rank = tensor.shape[1]
    per_head = tensor.reshape(NUM_HEADS, 3, HEAD_DIM, rank)
    return (
        per_head[:, 0].reshape(INNER_DIM, rank),
        per_head[:, 1].reshape(INNER_DIM, rank),
        per_head[:, 2].reshape(INNER_DIM, rank),
    )


def _convert_kohya_lora(state_dict) -> dict:
    """Turn a kohya / CivitAI MiniMax-H3 adapter into diffusers keys. Returns `{}` when nothing matched."""
    import torch

    modules = {}
    for key in state_dict:
        clean = key.replace(".default", "")
        for suffix, role in _LORA_SUFFIXES:
            if clean.endswith(suffix):
                modules.setdefault(clean[: -len(suffix)], {})[role] = key
                break

    converted, split_count, swap_count, skipped = {}, 0, 0, 0
    for module, roles in modules.items():
        if "down" not in roles or "up" not in roles:
            skipped += 1
            continue

        down = state_dict[roles["down"]]
        up = state_dict[roles["up"]]
        rank = down.shape[0]

        # alpha carries the scale: PEFT applies `alpha / rank`, so it is folded into lora_B here and the key dropped.
        if "alpha" in roles and rank:
            try:
                alpha = float(state_dict[roles["alpha"]].reshape(-1)[0])
                if alpha > 0 and abs(alpha - rank) > 1e-6:
                    up = up * (alpha / rank)
            except Exception:  # noqa: BLE001
                pass

        name = _dotted_module(_strip_lora_prefixes(module))
        if name is None:
            skipped += 1
            continue
        name = _rename_module(name)

        if name.endswith(".attn.qkv_proj"):
            if up.shape[0] != 3 * INNER_DIM:
                skipped += 1
                continue
            stem = name[: -len("qkv_proj")]
            for part, tensor in zip(("to_q", "to_k", "to_v"), _split_qkv_by_head(up)):
                converted[f"{stem}{part}.lora_A.weight"] = down
                converted[f"{stem}{part}.lora_B.weight"] = tensor.contiguous()
            split_count += 1
            continue

        if name.endswith(".mlp.fc1"):
            name = name[: -len(".mlp.fc1")] + ".ff.net.0.proj"
            half = up.shape[0] // 2
            up = torch.cat([up[half:], up[:half]], dim=0)
            swap_count += 1

        converted[f"{name}.lora_A.weight"] = down
        converted[f"{name}.lora_B.weight"] = up

    if converted:
        print(f"[lora] kohya conversion: {len(converted) // 2} layers, qkv split {split_count}, "
              f"fc1 swapped {swap_count}, skipped {skipped}")
    return converted


def _load_lora_state_dict(path: str) -> dict:
    """Read a `.safetensors` or a torch `.pt`/`.bin`, whichever the link handed over."""
    if path.lower().endswith((".pt", ".pth", ".bin", ".ckpt")):
        import torch

        loaded = torch.load(path, map_location="cpu", weights_only=True)
        for wrapper in ("state_dict", "lora", "module", "weights"):
            if isinstance(loaded, dict) and isinstance(loaded.get(wrapper), dict):
                loaded = loaded[wrapper]
                break
        return loaded

    from safetensors.torch import load_file

    return load_file(path)


def _is_comfyui_lora(state_dict) -> bool:
    """Whether a LoRA state dict is in ComfyUI's MiniMax-H3 naming rather than diffusers'.

    ComfyUI names the block stack `blocks.N.*` and the token refiner `token_refiner.blocks.N.*`; diffusers names them
    `transformer_blocks.N.*` and `token_refiner.refiner_blocks.N.*`. A key starting with `blocks.` is the tell.
    """
    for key in state_dict:
        if key.startswith(("blocks.", "token_refiner.blocks.", "final_layer.")):
            return True
    return False


def _convert_comfyui_lora(state_dict) -> dict:
    """Remap a ComfyUI-format MiniMax-H3 Turbo LoRA to the diffusers `transformer_ref` module names.

    The Turbo LoRA ([`larryvrh/MiniMax-H3-Turbo-Lora`](https://huggingface.co/larryvrh/MiniMax-H3-Turbo-Lora)) is trained
    against the ComfyUI checkpoint, whose module names differ from diffusers' in four ways:

    * the block stack is `blocks.N` in ComfyUI but `transformer_blocks.N` in diffusers,
    * the token refiner is `token_refiner.blocks.N` but `token_refiner.refiner_blocks.N`,
    * the final AdaLN is `final_layer.adaln_proj.linear` but `norm_out.linear`,
    * attention QKV is one fused `attn.qkv_proj` in ComfyUI but three separate `attn.to_q` / `to_k` / `to_v` in
      diffusers, and the output projection is `attn.out_proj` but `attn.to_out.0`,
    * the feed-forward is `mlp.fc1` / `mlp.fc2` but `ff.fc1` / `ff.fc2`.

    The fused QKV `lora_B` is `[3 * inner_dim, rank]`; splitting it into three along dim 0 gives the three separate
    `lora_B` matrices, and `lora_A` (which is `[rank, hidden_size]`) is shared verbatim across the three. The metadata
    says `W_eff = W + lora_B @ lora_A` with alpha = rank, so the scaling is 1.0 and no alpha key is added.
    """
    import torch

    converted = {}
    for key, value in state_dict.items():
        # `blocks.N.` -> `transformer_blocks.N.`
        if key.startswith("blocks."):
            new_key = "transformer_blocks." + key[len("blocks."):]
        elif key.startswith("token_refiner.blocks."):
            new_key = "token_refiner.refiner_blocks." + key[len("token_refiner.blocks."):]
        elif key.startswith("final_layer.adaln_proj.linear."):
            new_key = "norm_out.linear." + key[len("final_layer.adaln_proj.linear."):]
        else:
            converted[key] = value
            continue

        # At this point `new_key` is a diffusers block path. Remap the leaf module names.
        if ".attn.qkv_proj." in new_key:
            # Fused QKV: split `lora_B` along dim 0 into q/k/v, duplicate `lora_A` verbatim.
            leaf = new_key.split(".attn.qkv_proj.")[-1]  # `lora_A.weight` or `lora_B.weight`
            stem = new_key[: new_key.index(".attn.qkv_proj.")]
            if leaf == "lora_A.weight":
                for proj in ("to_q", "to_k", "to_v"):
                    converted[f"{stem}.attn.{proj}.lora_A.weight"] = value
            else:  # lora_B.weight
                q_b, k_b, v_b = value.chunk(3, dim=0)
                converted[f"{stem}.attn.to_q.lora_B.weight"] = q_b
                converted[f"{stem}.attn.to_k.lora_B.weight"] = k_b
                converted[f"{stem}.attn.to_v.lora_B.weight"] = v_b
        elif ".attn.out_proj." in new_key:
            converted[new_key.replace(".attn.out_proj.", ".attn.to_out.0.")] = value
        elif ".mlp." in new_key:
            converted[new_key.replace(".mlp.", ".ff.")] = value
        else:
            # `adaln_proj.linear` and the token refiner's attention/ff already match diffusers' names after the
            # block-prefix rename above.
            converted[new_key] = value
    return converted


def apply_loras(transformer, loras) -> list[str]:
    """Attach `loras` (local path, strength) to `transformer` and give each its strength, replacing whatever was on it.

    Every adapter already on the model is removed first, so a request is never affected by the one before it — which
    matters when a worker is reused rather than forked fresh. A LoRA in ComfyUI's MiniMax-H3 naming is remapped to
    diffusers' module names on the fly, so the Turbo LoRA works without a separate conversion step.
    """
    import torch

    from safetensors.torch import load_file

    for name in list(getattr(transformer, "peft_config", None) or {}):
        transformer.delete_adapters(name)

    names, scales = [], []
    for index, (path, scale) in enumerate(loras):
        state_dict = _load_lora_state_dict(path)
        if not any(marker in key for key in state_dict for marker in _DIFFUSERS_MARKERS):
            # kohya / CivitAI naming first, since it also covers the flat underscored form; the older ComfyUI
            # remap stays as the fallback for files the kohya pass does not recognise.
            converted = _convert_kohya_lora(state_dict)
            state_dict = converted or (
                _convert_comfyui_lora(state_dict) if _is_comfyui_lora(state_dict) else state_dict
            )
        name = f"lora{index}"
        transformer.load_lora_adapter(state_dict, adapter_name=name, prefix=_lora_prefix(state_dict))
        names.append(name)
        scales.append(float(scale))

    if not names:
        return []

    # PEFT builds the new layers on its own default device/dtype; the base weights are the truth here, under either
    # placement mode (`offload` keeps them on the host and moves whole modules by hook).
    base = next(param for key, param in transformer.named_parameters() if ".lora_" not in key)
    with torch.no_grad():
        for key, param in transformer.named_parameters():
            if ".lora_" in key and (param.device != base.device or param.dtype != base.dtype):
                param.data = param.data.to(device=base.device, dtype=base.dtype)

    transformer.set_adapters(names, scales)
    return names


def collect_loras(lora_fields, progress) -> tuple[list[tuple[str, float]], list[str]]:
    """Resolve the UI's `reference, strength, reference, strength, ...` into `(local path, strength)` pairs.

    Resolved before the booking: a download that happens inside `@spaces.GPU` is billed as GPU time.
    """
    loras, labels = [], []
    for reference, scale in zip(lora_fields[::2], lora_fields[1::2]):
        reference = (reference or "").strip()
        if not reference or abs(float(scale)) < 1e-6:
            continue
        progress(0.0, desc=f"Fetching LoRA {reference} ...")
        loras.append((resolve_lora(reference), float(scale)))
        labels.append(f"{os.path.basename(reference)} @ {float(scale):g}")
    if loras and os.environ.get("H3_AOTI") == "1":
        raise gr.Error("A LoRA cannot be applied to an AoTI-compiled transformer. Turn `H3_AOTI` off.")
    return loras, labels


@cache
def conditioner():
    """The other half, over the gradio API. `gradio_client` attaches the caller's own ZeroGPU token per call, so the
    conditioner's booking is billed to whoever asked for the video."""
    from gradio_client import Client

    return Client(CONDITIONER_SPACE)


def probe(path: str) -> tuple[float | None, float | None]:
    """`(video seconds, audio seconds)` of a media file, either being `None` when the stream is absent."""
    import av

    def seconds(stream, container):
        if stream.duration is not None and stream.time_base is not None:
            return float(stream.duration * stream.time_base)
        return None if container.duration is None else container.duration / av.time_base

    with av.open(path) as container:
        video = seconds(container.streams.video[0], container) if container.streams.video else None
        audio = seconds(container.streams.audio[0], container) if container.streams.audio else None
    return video, audio


def collect(image_paths, audio_path, video_path) -> list[tuple[str, str]]:
    """The `(kind, path)` references of a request, **in the order the model reads them**.

    That order numbers the labels of MiniMax-H3's prompt presentation and advances the shared audio/video rotary clock,
    so the same references in a different order are a different request.
    """
    ordered = [("image", path) for path in image_paths if path]
    if audio_path:
        ordered.append(("audio", audio_path))
    if video_path:
        ordered.append(("video", video_path))
    return ordered


def build_references(references: list[tuple[str, str]]):
    """The `(kind, path)` references of a request as decoded reference dataclasses, in packed order. `from_file` brings
    the rates along: a video its own frame rate and soundtrack, a clip its sample rate."""
    from diffusers.modular_pipelines.minimax_h3 import (
        MiniMaxH3AudioReference,
        MiniMaxH3ImageReference,
        MiniMaxH3VideoReference,
    )

    classes = {"image": MiniMaxH3ImageReference, "video": MiniMaxH3VideoReference, "audio": MiniMaxH3AudioReference}
    return [classes[kind].from_file(path) for kind, path in references]


def audio_bearing(references: list[tuple[str, str]]) -> list[tuple[str, float]]:
    """The references that carry a waveform, and how long it is. A video reference brings its own soundtrack."""
    carried = []
    for kind, path in references:
        if kind == "image":
            continue
        _, audio_seconds = probe(path)
        if audio_seconds is not None:
            carried.append((kind, audio_seconds))
    return carried


def duration_controls(audio_path, video_path, match: bool):
    """Show the duration slider unless a single soundtrack can set it, which is when MiniMax-H3 lets it be left out."""
    try:
        carried = audio_bearing(collect([], audio_path, video_path))
    except Exception:
        carried = []
    # Exactly one soundtrack, long enough to be a duration MiniMax-H3 generates; anything else is ambiguous or out of
    # range and the slider stays.
    derivable = len(carried) == 1 and MIN_DURATION <= snap_frames(carried[0][1]) / FPS <= MAX_REFERENCE_VIDEO
    return gr.update(visible=derivable), gr.update(visible=not (derivable and match))


def check(prompt: str, references: list[tuple[str, str]]) -> None:
    """The model's own rules, before anything is uploaded or a card is allocated."""
    if not prompt or not prompt.strip():
        raise gr.Error("MiniMax-H3 always takes a prompt, references or not.")
    if not references:
        raise gr.Error("Add at least one reference — an image or a video for the model to condition on.")
    if {kind for kind, _ in references} == {"audio"}:
        raise gr.Error("An audio reference needs an image or a video alongside it; it cannot go on its own.")
    for kind, path in references:
        if kind != "video":
            continue
        video_seconds, _ = probe(path)
        if video_seconds is None:
            raise gr.Error("That reference video has no video stream. Drop it in the audio slot instead.")
        if not MIN_REFERENCE_VIDEO <= video_seconds <= MAX_REFERENCE_VIDEO:
            raise gr.Error(
                f"The reference video is {video_seconds:.1f} s. Use a clip between "
                f"{MIN_REFERENCE_VIDEO:g} and {MAX_REFERENCE_VIDEO:g} seconds."
            )


def encode_remote(prompt, references, canvas, num_frames, rewrite_prompt=False):
    """`/encode_ref2va` on the conditioner Space: a safetensors file holding `prompt_embeds` + `text_token_tags`, with
    the resolved `height` / `width` / `num_frames` in its metadata, plus the plan.

    `canvas` is the label. `media` and `kinds` are parallel and ordered, and the references go over because `ref2va`'s
    presentation puts a vision block in front of the prompt for every image and every merged video frame pair.
    """
    from gradio_client import handle_file
    from safetensors import safe_open

    path, plan = conditioner().predict(
        prompt=prompt,
        media=[handle_file(path) for _, path in references],
        kinds=",".join(kind for kind, _ in references),
        canvas=canvas,
        num_frames=num_frames,
        rewrite_prompt=bool(rewrite_prompt),
        api_name="/encode_ref2va",
    )
    with safe_open(path, framework="pt") as handle:
        return handle.get_tensor("prompt_embeds"), handle.get_tensor("text_token_tags"), handle.metadata(), plan


@spaces.GPU(duration=get_duration, size=GPU_SIZE)
def _generate(prompt_embeds, text_token_tags, references, height, width, num_frames, steps, seed, loras=()):
    """The only thing on GPU time: the two reference encoders, the packed-sequence denoise loop and the decoders.

    References cross as paths and are decoded here; only the three generated outputs come back. A `@spaces.GPU`
    argument crosses a process boundary by pickling, a 5 s 1344x768 reference video is 370 MB of expanded frames, and
    the full `PipelineState` still holds the packed latents and the rotary grid on the card.

    The adapters are attached here rather than in the caller: `spaces` runs this body in its own worker, so the
    transformer the request sees is the one that has to carry them.
    """
    import torch

    if PLACEMENT == "lazy":
        PIPE.to("cuda")

    apply_loras(PIPE.transformer_ref, loras or ())

    state = PIPE(
        prompt_embeds=prompt_embeds.to("cuda"),
        text_token_tags=text_token_tags,
        references=build_references(references),
        height=height,
        width=width,
        num_frames=num_frames,
        num_inference_steps=int(steps),
        generator=torch.Generator("cpu").manual_seed(int(seed)),
    )
    return state.get("videos")[0], state.get("audio")[0].cpu(), state.get("sampling_rate")


def generate(
    # Every parameter after `prompt` has a default, and the newest ones sit at the end, so a positional API client
    # written against an older signature keeps working.
    prompt,
    image_1=None,
    audio_path=None,
    video_path=None,
    canvas=DEFAULT_CANVAS,
    image_2=None,
    image_3=None,
    image_4=None,
    image_5=None,
    image_6=None,
    image_7=None,
    image_8=None,
    image_9=None,
    match=True,
    duration=5,
    steps=28,
    seed=42,
    upsample=False,
    *lora_fields,
    progress=gr.Progress(track_tqdm=True),
):
    """One request. The LoRA fields are last and default to empty, so a positional API client that predates them is
    unaffected. `lora_fields` arrives as `reference, strength, reference, strength, ...`."""
    if LOAD_ERROR:
        raise gr.Error(LOAD_ERROR)
    if PIPE is None:
        raise gr.Error("The denoiser is still loading.")

    from diffusers.utils import encode_video

    images = [image_1, image_2, image_3, image_4, image_5, image_6, image_7, image_8, image_9]
    references = collect(images, audio_path, video_path)
    check(prompt, references)

    # `0` is "leave it to the references" over the wire, which MiniMax-H3 accepts when exactly one of them carries a
    # soundtrack. The conditioner resolves it either way and this Space pins whatever comes back.
    derivable = len(audio_bearing(references)) == 1
    requested = 0 if (match and derivable) else snap_frames(duration)

    loras, lora_labels = collect_loras(lora_fields, progress)

    # Before the conditioner spends GPU time: the canvas from the table, worst case for the frame count.
    planned_height, planned_width = CANVASES.get(canvas, CANVASES[DEFAULT_CANVAS])
    fits(
        TEXT_TOKEN_ALLOWANCE,
        references,
        planned_height,
        planned_width,
        requested or snap_frames(MAX_REFERENCE_VIDEO),
        steps,
        loras,
    )

    progress(0.0, desc="Upsampling the prompt ..." if upsample else "Reading the prompt and references ...")
    conditioned = time.time()
    try:
        prompt_embeds, text_token_tags, metadata, plan = encode_remote(
            prompt, references, canvas, requested, rewrite_prompt=upsample
        )
    except gr.Error:
        raise
    except Exception as error:
        # gradio only puts the exception *type* on the wire, so the useful half of a conditioner-side failure is in
        # that Space's logs.
        traceback.print_exc()
        raise gr.Error(
            f"The conditioner ({CONDITIONER_SPACE}) failed with `{type(error).__name__}: {error}`. "
            "Its logs carry the full traceback."
        ) from error
    condition_seconds = time.time() - conditioned
    height, width, num_frames = (int(metadata[key]) for key in ("height", "width", "num_frames"))
    refined = plan.get("refined_prompt") or ""

    # Again, with the exact numbers the conditioner returned.
    fits(int(text_token_tags.shape[0]), references, height, width, num_frames, steps, loras)

    progress(0.1, desc=f"Generating {num_frames / FPS:.1f} s at {width}x{height} ...")
    started = time.time()
    frames, audio, sampling_rate = _generate(
        prompt_embeds, text_token_tags, references, height, width, num_frames, steps, seed, loras
    )
    generate_seconds = time.time() - started

    directory = os.path.join(tempfile.gettempdir(), "h3-outputs")
    os.makedirs(directory, exist_ok=True)
    path = os.path.join(directory, f"h3-ref2va-{int(time.time() * 1000)}.mp4")
    encode_video(frames, fps=FPS, output_path=path, audio=audio, audio_sample_rate=sampling_rate)

    print(
        f"[ref2va] {[kind for kind, _ in references]} · `{width}x{height}`, {num_frames} frames "
        f"({num_frames / FPS:.3f} s), {int(steps)} steps · conditioner {condition_seconds:.0f}s "
        f"({plan['num_text_tokens']} tokens{', upsampled' if refined else ''}) · "
        f"denoise + decode {generate_seconds:.0f}s "
        f"({generate_seconds / int(steps):.1f} s/step) · seed {int(seed)}"
        f"{' · LoRA ' + ', '.join(lora_labels) if lora_labels else ''}",
        flush=True,
    )
    return path, refined, gr.update(visible=bool(refined))


# ----------------------------------------------------------------------------------------------------------------
# Settings file
# ----------------------------------------------------------------------------------------------------------------
# Everything typed rather than uploaded, so a session can be picked up where it was left off. The references
# themselves are deliberately left out: gradio hands them over as paths into a per-session temporary directory that
# is gone by the next visit, so a saved path would restore as a dead file rather than as the image.

SETTINGS_VERSION = 1
SETTINGS_KEYS = (
    ["prompt", "upsample", "canvas", "match", "duration", "steps", "seed"]
    + [f"lora_{slot + 1}" for slot in range(LORA_SLOTS)]
    + [f"lora_{slot + 1}_scale" for slot in range(LORA_SLOTS)]
    + ["randomize_seed"]
)

MAX_SEED = 2**31 - 1


def roll_seed(randomize, seed):
    """A fresh seed per press when the box is ticked, otherwise the one that is set."""
    return random.randint(0, MAX_SEED) if randomize else int(seed)


def save_settings(*values):
    """Write the current controls to a `.json` and reveal it for download."""
    payload = {"version": SETTINGS_VERSION, "saved": time.strftime("%Y-%m-%d %H:%M:%S")}
    payload.update(dict(zip(SETTINGS_KEYS, values)))

    directory = os.path.join(tempfile.gettempdir(), "h3-settings")
    os.makedirs(directory, exist_ok=True)
    path = os.path.join(directory, f"h3-settings-{int(time.time())}.json")
    with open(path, "w", encoding="utf-8") as handle:
        json.dump(payload, handle, ensure_ascii=False, indent=2, default=str)
    return gr.update(value=path, visible=True)


def load_settings(path):
    """Restore the controls from a `.json`. A key the file does not carry leaves its control alone, so a settings
    file written by an older version of this Space still loads."""
    if not path:
        return [gr.update() for _ in SETTINGS_KEYS]
    try:
        with open(path, encoding="utf-8") as handle:
            payload = json.load(handle)
    except Exception as error:
        raise gr.Error(f"That settings file cannot be read: `{type(error).__name__}: {error}`")
    if not isinstance(payload, dict):
        raise gr.Error("That is not a settings file for this Space.")

    updates = []
    for key in SETTINGS_KEYS:
        value = payload.get(key)
        # An unknown canvas label would be rejected by the conditioner, which is the wrong place to find out.
        if value is None or (key == "canvas" and value not in CANVASES):
            updates.append(gr.update())
        else:
            updates.append(gr.update(value=value))
    return updates


# ----------------------------------------------------------------------------------------------------------------
# Named profiles
# ----------------------------------------------------------------------------------------------------------------
# Same payload as the settings file, but stored under a name inside the Space, so a set-up can be recalled from a
# dropdown instead of a download / re-upload round trip. `/data` is used when the Space has persistent storage
# attached, so the profiles survive a restart; otherwise they live for as long as the Space runs.

if os.path.isdir("/data") and os.access("/data", os.W_OK):
    PROFILES_DIR = os.path.join("/data", "h3-profiles")
else:
    PROFILES_DIR = os.path.join(tempfile.gettempdir(), "h3-profiles")
os.makedirs(PROFILES_DIR, exist_ok=True)

NO_PROFILE = "— no saved profile —"


def _profile_file(name: str) -> str:
    safe = re.sub(r"[^A-Za-z0-9 ._-]", "_", (name or "").strip())[:64].strip() or "profile"
    return os.path.join(PROFILES_DIR, f"{safe}.json")


def list_profiles() -> list[str]:
    names = []
    for entry in sorted(os.listdir(PROFILES_DIR)) if os.path.isdir(PROFILES_DIR) else []:
        if entry.endswith(".json"):
            names.append(entry[:-5])
    return names


def save_profile(name, *values):
    """Store the current controls under a name and reselect it in the dropdown."""
    if not (name or "").strip():
        return gr.update(), "Give the profile a name first."
    payload = {"version": SETTINGS_VERSION, "saved": time.strftime("%Y-%m-%d %H:%M:%S")}
    payload.update(dict(zip(SETTINGS_KEYS, values)))
    path = _profile_file(name)
    try:
        with open(path, "w", encoding="utf-8") as handle:
            json.dump(payload, handle, ensure_ascii=False, indent=2, default=str)
    except Exception as error:
        return gr.update(), f"Could not save: `{type(error).__name__}: {error}`"
    saved = os.path.basename(path)[:-5]
    return gr.update(choices=[NO_PROFILE, *list_profiles()], value=saved), f"Saved **{saved}**."


def load_profile(name):
    """Restore every control from a named profile."""
    blank = [gr.update() for _ in SETTINGS_KEYS]
    if not name or name == NO_PROFILE:
        return [*blank, gr.update(), ""]
    path = _profile_file(name)
    if not os.path.exists(path):
        return [*blank, gr.update(), f"No profile named **{name}**."]
    try:
        with open(path, encoding="utf-8") as handle:
            payload = json.load(handle)
    except Exception as error:
        return [*blank, gr.update(), f"Could not read it: `{type(error).__name__}: {error}`"]

    updates = []
    for key in SETTINGS_KEYS:
        value = payload.get(key)
        if value is None or (key == "canvas" and value not in CANVASES):
            updates.append(gr.update())
        else:
            updates.append(gr.update(value=value))
    stamp = payload.get("saved", "")
    return [*updates, gr.update(value=name), f"Loaded **{name}**{f' (saved {stamp})' if stamp else ''}."]


def refresh_profiles(current=None):
    """Re-read the folder. The dropdown's choices are built once at start-up, so a profile saved in
    another tab (or after this page was opened) would otherwise stay invisible until a restart."""
    names = list_profiles()
    value = current if current in names else NO_PROFILE
    return gr.update(choices=[NO_PROFILE, *names], value=value)


def delete_profile(name):
    if not name or name == NO_PROFILE:
        return gr.update(), "Pick a profile first."
    path = _profile_file(name)
    try:
        os.remove(path)
    except FileNotFoundError:
        return gr.update(choices=[NO_PROFILE, *list_profiles()], value=NO_PROFILE), f"No profile named **{name}**."
    except Exception as error:
        return gr.update(), f"Could not delete: `{type(error).__name__}: {error}`"
    return gr.update(choices=[NO_PROFILE, *list_profiles()], value=NO_PROFILE), f"Deleted **{name}**."


# ----------------------------------------------------------------------------------------------------------------
# Live GPU cost
# ----------------------------------------------------------------------------------------------------------------


def gpu_estimate(canvas, duration, steps, match, audio_path, video_path, *rest):
    """What this request will reserve, by the same `budget()` the pre-flight check and `get_duration` use."""
    images = list(rest[:MAX_IMAGE_SLOTS])
    lora_fields = [value for value in rest[MAX_IMAGE_SLOTS:] if (value or "").strip()]
    try:
        references = collect(images, audio_path, video_path)
    except Exception:
        references = []

    try:
        derivable = len(audio_bearing(references)) == 1
    except Exception:
        derivable = False

    try:
        num_frames = snap_frames(float(duration))
        height, width = CANVASES.get(canvas, CANVASES[DEFAULT_CANVAS])
        sequence, total, per_step, overhead = budget(
            TEXT_TOKEN_ALLOWANCE, references, height, width, num_frames, int(steps), lora_fields
        )
    except Exception as error:
        return f"⏳ **GPU cost:** estimate unavailable (`{type(error).__name__}`)"

    seconds = num_frames / FPS
    reserved = max(MIN_GPU_DURATION, min(MAX_GPU_DURATION, int(total)))
    room = int((MAX_GPU_DURATION - overhead) / per_step) if per_step > 0 else 0

    if sequence > MAX_SEQUENCE:
        head = (
            f"🚫 **Too large for the card:** {sequence} rows against a ceiling of {MAX_SEQUENCE}. "
            "Lower the duration, pick a smaller canvas, or remove a reference."
        )
    elif total > MAX_GPU_DURATION:
        advice = f"Lower Steps to {room}." if room >= MIN_STEPS else "Lower the duration or pick a smaller canvas."
        head = f"🚫 **Wants ~{int(total)} s of GPU, ceiling is {MAX_GPU_DURATION} s.** {advice}"
    else:
        head = f"⏳ **GPU cost: ~{reserved} s**"

    detail = (
        f"{width}x{height} · {seconds:.1f} s ({num_frames} frames) · {int(steps)} steps · "
        f"{len(references)} reference(s) · {len(lora_fields)} LoRA · {sequence} rows · "
        f"~{per_step:.1f} s per step"
    )
    if derivable and match:
        detail += " · duration comes from the reference soundtrack"
    return f"{head}  \n<sub>{detail}</sub>"


def _fill_lora_slots(files, *current):
    """Drop `.safetensors` files on the uploader and their paths land in the first free slots, so a local adapter
    needs no typing at all."""
    slots = list(current)
    for path in files or []:
        for index, value in enumerate(slots):
            if not (value or "").strip():
                slots[index] = path
                break
    return [gr.update(value=value) for value in slots]


def _add_preset_lora(preset, *current):
    """Fill the first free LoRA slot with a preset adapter, set its strength to 1.0 and move the steps slider to the
    preset's recommended count.

    The Turbo presets are tuned for a specific step range, so the steps slider is moved along with the slot — it is the
    one output beyond the LoRA fields. A slot already holding the same reference is a no-op, so the button can be
    pressed twice without duplicating, and a full set of slots is left untouched.
    """
    reference, steps, _, strength = LORA_PRESETS[preset]
    slots = list(current[:LORA_SLOTS])
    scales = list(current[LORA_SLOTS:])
    if reference not in [(value or "").strip() for value in slots]:
        for index, value in enumerate(slots):
            if not (value or "").strip():
                slots[index] = reference
                scales[index] = strength
                break
    return [*slots, *scales, steps]


load_models()

CHIPS = [
    "cinematic lighting, shallow depth of field",
    "slow push in",
    "camera orbits the subject",
    "handheld camera, subtle shake",
    "the character speaks to the camera",
    "rain, neon reflections",
    "warm golden hour light",
    "ambient room tone, quiet footsteps",
]

# ---------------------------------------------------------------------------------------------
# Clip stitching - CPU only, zero GPU quota
# ---------------------------------------------------------------------------------------------
# Runs on videos that already exist, so it never touches the card. Three 5 s clips become one 15 s
# file for no extra ZeroGPU time. LTX-2.5 carries a soundtrack, so the audio track is carried
# through too: every clip is normalised to one frame size and one audio format first, and a clip
# that somehow has no audio gets silence rather than breaking the join.


def _ffmpeg_exe() -> str:
    """imageio-ffmpeg ships a binary, so this works even without a system ffmpeg."""
    try:
        import imageio_ffmpeg
        return imageio_ffmpeg.get_ffmpeg_exe()
    except Exception:  # noqa: BLE001
        return "ffmpeg"


def _probe(path):
    """(width, height, has_audio), read out of ffmpeg's own report on the file."""
    result = subprocess.run([_ffmpeg_exe(), "-i", path], capture_output=True)
    text = result.stderr.decode("utf-8", "ignore")
    has_audio = "Audio:" in text
    size = re.search(r"Video:.*?(\d{2,5})x(\d{2,5})", text)
    if size:
        return int(size.group(1)), int(size.group(2)), has_audio
    return 0, 0, has_audio


def _normalise(path, width, height, out_path):
    scale = (f"scale={width}:{height}:force_original_aspect_ratio=decrease,"
             f"pad={width}:{height}:(ow-iw)/2:(oh-ih)/2,setsar=1,fps={FPS}")
    _, _, has_audio = _probe(path)
    if has_audio:
        command = [_ffmpeg_exe(), "-y", "-i", path, "-vf", scale,
                   "-c:v", "libx264", "-preset", "fast", "-crf", "20",
                   "-c:a", "aac", "-ar", "48000", "-ac", "2", out_path]
    else:
        command = [_ffmpeg_exe(), "-y", "-i", path,
                   "-f", "lavfi", "-i", "anullsrc=r=48000:cl=stereo",
                   "-vf", scale, "-c:v", "libx264", "-preset", "fast", "-crf", "20",
                   "-c:a", "aac", "-ar", "48000", "-ac", "2", "-shortest",
                   "-map", "0:v:0", "-map", "1:a:0", out_path]
    subprocess.run(command, check=True, capture_output=True)
    return out_path


def concat_videos(paths, name_hint=""):
    valid = [p for p in (paths or []) if p and os.path.exists(p)]
    if not valid:
        return None
    if len(valid) == 1:
        return valid[0]

    width, height, _ = _probe(valid[0])
    width, height = width or 960, height or 544
    work_dir = tempfile.mkdtemp(prefix="stitch-")
    clean = str(name_hint).strip().replace(" ", "_") or f"stitched_{len(valid)}clips"
    if clean.lower().endswith(".mp4"):
        clean = clean[:-4]
    out_path = os.path.join(work_dir, f"{clean}.mp4")

    try:
        parts = [
            _normalise(path, width, height, os.path.join(work_dir, f"part{index}.mp4"))
            for index, path in enumerate(valid)
        ]
        inputs = []
        for part in parts:
            inputs += ["-i", part]
        chain = "".join(f"[{i}:v][{i}:a]" for i in range(len(parts)))
        filter_str = f"{chain}concat=n={len(parts)}:v=1:a=1[outv][outa]"
        subprocess.run(
            [_ffmpeg_exe(), "-y", *inputs, "-filter_complex", filter_str,
             "-map", "[outv]", "-map", "[outa]",
             "-c:v", "libx264", "-preset", "fast", "-crf", "20",
             "-c:a", "aac", "-movflags", "+faststart", out_path],
            check=True, capture_output=True,
        )
    except Exception as error:  # noqa: BLE001
        raise gr.Error(f"Stitching failed: {error}")
    return out_path


def _queue_status(count, path=None):
    if not count:
        return "Queue: empty."
    name = f" → `{os.path.basename(path)}`" if path else ""
    return f"**{count} clip(s) in the queue**{name}"


def add_to_queue(video_path, queue, name_hint):
    queue = list(queue or [])
    if video_path and os.path.exists(str(video_path)):
        queue.append(video_path)
    if not queue:
        return None, None, [], "Nothing to add yet — generate a video first."
    merged = concat_videos(queue, name_hint)
    return merged, merged, queue, _queue_status(len(queue), merged)


def clear_queue():
    return None, None, [], _queue_status(0)


def auto_queue(video_path, auto_on, queue, name_hint):
    """Runs after every generation; only appends when the box is ticked."""
    queue = list(queue or [])
    if not auto_on:
        return gr.update(), gr.update(), queue, _queue_status(len(queue))
    return add_to_queue(video_path, queue, name_hint)


CSS = """
@import url('https://fonts.googleapis.com/css2?family=Inter:wght@400;500;600;700&display=swap');

.gradio-container { font-family: 'Inter', ui-sans-serif, system-ui, sans-serif !important; }
.main.fillable { max-width: 1400px !important; }
.dark .gradio-container { color: var(--body-text-color); }

#h3-hero {
    border-radius: 18px;
    padding: 22px 26px;
    margin-bottom: 12px;
    background: linear-gradient(120deg, #0b1120 0%, #1d4ed8 45%, #0891b2 100%);
    color: #fff;
    box-shadow: 0 10px 30px rgba(8, 47, 73, .28);
}
#h3-hero h1 { margin: 0; font-size: 25px; font-weight: 700; letter-spacing: -.02em; color: #fff; }
#h3-hero p { margin: 6px 0 0 0; font-size: 14px; opacity: .92; color: #fff; }
#h3-hero a { color: #fff; text-decoration: underline; }
#h3-hero .pills { margin-top: 12px; display: flex; flex-wrap: wrap; gap: 8px; }
#h3-hero .pills span {
    background: rgba(255,255,255,.16);
    border: 1px solid rgba(255,255,255,.25);
    padding: 4px 11px; border-radius: 999px; font-size: 12px; font-weight: 500;
}

.panel {
    border-radius: 16px !important;
    border: 1px solid var(--border-color-primary) !important;
    padding: 14px !important;
    background: var(--background-fill-secondary) !important;
}

#run-btn {
    font-size: 17px !important;
    font-weight: 700 !important;
    min-height: 62px !important;
    border-radius: 14px !important;
    box-shadow: 0 8px 22px rgba(29, 78, 216, .25);
}
#profile-refresh, #seed-dice, #turbo-btn, #profile-save, #profile-load, #profile-delete {
    min-height: 42px !important; border-radius: 12px !important; font-weight: 600 !important;
}

.chip-row { gap: 6px !important; }
.chip-row button {
    border-radius: 999px !important;
    font-size: 12.5px !important;
    padding: 6px 12px !important;
    min-height: 34px !important;
    font-weight: 500 !important;
}

#estimate-btn { min-height: 44px !important; border-radius: 12px !important; }
.gpu-estimate {
    border-radius: 12px;
    padding: 10px 14px;
    background: rgba(37, 99, 235, .10);
    border: 1px solid rgba(37, 99, 235, .35);
    font-size: 14px;
}
.gpu-estimate sub { opacity: .75; font-size: 12px; }

.turbo-blurb { font-size: 13px; opacity: .85; }
footer { display: none !important; }
"""

THEME = gr.themes.Soft(primary_hue="blue", secondary_hue="cyan", neutral_hue="slate", radius_size="lg")

HERO = """
<div id="h3-hero">
  <h1>MiniMax-H3 &middot; reference &rarr; video + soundtrack &middot; Custom lora + CivitAI links, GPU cost, profiles, clip stitching</h1>
  <p>33B model generating video and a fully synchronized soundtrack (ambience, foley, speech) from your own subject,
  voice or camera move.
  <a href="https://huggingface.co/MiniMaxAI/MiniMax-H3" target="_blank" rel="noopener">model</a> &middot;
  <a href="https://www.minimax.io/blog/minimax-h3" target="_blank" rel="noopener">blog</a> &middot;
  <a href="https://huggingface.co/spaces/multimodalart/minimax-h3" target="_blank" rel="noopener">text / image to video</a></p>
  <div class="pills">
    <span>33B</span><span>joint video + audio</span><span>Turbo lora: 4&ndash;8 steps</span>
    <span>ComfyUI lora accepted</span><span>up to 9 references</span><span>5 custom lora slots</span><span>named profiles</span><span>CivitAI links</span><span>kohya auto-convert</span><span>GPU cost estimate</span><span>clip stitching with audio</span>
  </div>
</div>
"""

LORA_HELP = """Each slot takes a Hugging Face repo (`owner/repo`), a file inside one
(`owner/repo/name.safetensors`), a **CivitAI download link**, any other direct `.safetensors` URL, or a local path —
or just drop the files below. A strength of `0` switches a slot off without clearing it. Slots start at **0.5**,
which is what most H3 adapters on CivitAI are written up for — raise it if the effect is too weak, and keep the sum
of several adapters in mind, since stacking three at `1.0` is what turns a clip plastic. Adapters have to be trained
against the `transformer_ref/` partition. ComfyUI-trained adapters are remapped to diffusers' module names on the fly,
so no separate conversion step is needed.

CivitAI links are fetched by the Space itself, not by your browser, so being signed in there does not help a gated
model — set `CIVITAI_TOKEN` under *Settings → Variables and secrets* and it is appended automatically. Downloads are
cached, so a second run with the same link costs nothing.

**kohya / CivitAI files are converted on the fly.** Their flat underscored names are re-dotted, the fused QKV is
un-interleaved head by head (a plain three-way split hands q's rows to k), the gated MLP's two halves are put back in
diffusers' order, `alpha` is folded into the weights, and `.pt` files are read as well as `.safetensors`. Nothing has
to be converted by hand first.
"""

TURBO_HELP = """A few-step distillation that renders joint video + soundtrack in **4–8 steps** instead of the usual
~28 (a ~5× speedup). Pick one and press **add** — it fills a free slot at that build's own strength and moves the
steps slider to the count it was trained for.

Larryvrh documents strength `1.0` for every build, and that is what the **v4** entries are filled at. The older **v1**
line is the one people report over-sharpening and plastic skin on at `1.0`, so those are filled at `0.7` instead. Both
are starting points, not rules — if a clip smears, raise it; if it looks over-sharp or plastic, lower it.
"""

PROFILE_HELP = """Profiles keep the prompt, the canvas, the sliders and the lora slots — everything typed rather than
uploaded. Images, audio and video are not stored: gradio keeps them in a temporary folder that is gone by the next
visit, so a saved path would come back as a dead file.
"""

# gradio 6.0 takes `theme` and `css` on launch(), not here - passing them to the constructor
# only earns a warning and the styling is dropped.
with gr.Blocks(title="MiniMax-H3 - Custom lora + CivitAI, GPU cost, profiles, clip stitching") as demo:
    gr.HTML(HERO)

    with gr.Row(equal_height=False):
        with gr.Column(scale=5):

            # ---------------- prompt ----------------
            with gr.Group(elem_classes="panel"):
                prompt = gr.Textbox(
                    label="✍️ Prompt",
                    lines=3,
                    value="The character walks through a neon-lit street in the rain, humming to themselves",
                )
                upsample = gr.Checkbox(label="✨ Upsample prompt", value=False)
                with gr.Accordion("💡 Quick tags — click to add", open=False):
                    with gr.Row(elem_classes="chip-row"):
                        chip_buttons_a = [gr.Button(text, size="sm", variant="secondary") for text in CHIPS[:4]]
                    with gr.Row(elem_classes="chip-row"):
                        chip_buttons_b = [gr.Button(text, size="sm", variant="secondary") for text in CHIPS[4:]]

            # ---------------- references ----------------
            with gr.Group(elem_classes="panel"):
                # One tab per modality, in the order the model reads them. A reference left in a tab that is not the
                # open one is still part of the request.
                with gr.Tabs():
                    with gr.Tab("🖼️ Images"):
                        # One `gr.Row`, so gradio splits the width evenly and wraps at `min_width` rather than leaving
                        # a hole where a hidden slot used to be.
                        with gr.Row():
                            images = [
                                gr.Image(
                                    label="Subject, style or scene",
                                    type="filepath",
                                    min_width=180,
                                    # Fixed, so a row that wraps to a single slot stays the size of a full one.
                                    height=210,
                                    visible=index < OPEN_IMAGE_SLOTS,
                                )
                                for index in range(MAX_IMAGE_SLOTS)
                            ]
                        add_image = gr.Button("+ Add another image", size="sm", variant="secondary")
                    with gr.Tab("🔊 Audio"):
                        audio = gr.Audio(label="A voice or a piece of music", type="filepath")
                    with gr.Tab("🎥 Video"):
                        video = gr.Video(label="Motion & camera, 2–15 s. Its soundtrack comes along.")

            # ---------------- speed ----------------
            with gr.Group(elem_classes="panel"):
                gr.Markdown("### ⚡ Turbo lora (fast presets)")
                gr.Markdown(TURBO_HELP, elem_classes="turbo-blurb")
                with gr.Row():
                    lora_preset = gr.Dropdown(
                        label="Preset",
                        choices=list(LORA_PRESETS),
                        value=list(LORA_PRESETS)[0],
                        scale=4,
                    )
                    lora_preset_add = gr.Button("⚡ add to a free slot", variant="secondary", scale=1,
                                                elem_id="turbo-btn")
                turbo_blurb = gr.Markdown(
                    LORA_PRESETS[list(LORA_PRESETS)[0]][2], elem_classes="turbo-blurb"
                )

            # ---------------- the rest, in tabs ----------------
            with gr.Tabs():
                with gr.Tab(f"⭐ Custom lora ({LORA_SLOTS} slots)"):
                    gr.Markdown(LORA_HELP)
                    lora_references, lora_scales = [], []
                    for slot in range(LORA_SLOTS):
                        with gr.Row():
                            lora_references.append(
                                gr.Textbox(label=f"lora {slot + 1}", placeholder="owner/repo", scale=3)
                            )
                            lora_scales.append(
                                gr.Slider(
                                    label="Strength",
                                    minimum=LORA_MIN_SCALE,
                                    maximum=LORA_MAX_SCALE,
                                    step=0.05,
                                    # Most H3 adapters on CivitAI are written up for 0.5, so an empty slot starts
                                    # there rather than at 1.0. A Turbo preset overwrites it with its own number.
                                    value=DEFAULT_LORA_SCALE,
                                    scale=2,
                                )
                            )
                    lora_upload = gr.File(
                        label="Drop .safetensors here to fill the slots",
                        file_count="multiple",
                        file_types=[".safetensors"],
                        type="filepath",
                    )

                with gr.Tab("🎛️ Output"):
                    canvas = gr.Dropdown(label="Canvas", choices=list(CANVASES), value=DEFAULT_CANVAS)
                    match = gr.Checkbox(label="Match the reference soundtrack", value=True, visible=False)
                    duration = gr.Slider(
                        label="Duration (s)", minimum=MIN_DURATION, maximum=MAX_UI_DURATION, step=1, value=5
                    )
                    steps = gr.Slider(label="Steps", minimum=MIN_STEPS, maximum=40, step=1, value=28)
                    with gr.Row():
                        seed = gr.Number(label="Seed", value=42, precision=0, scale=3)
                        seed_dice = gr.Button("🎲 roll", variant="secondary", scale=1, elem_id="seed-dice")
                    randomize_seed = gr.Checkbox(
                        label="🎲 Randomize seed on every run",
                        value=True,
                        info="A new seed is drawn each time Generate is pressed, and lands in the box above.",
                    )

                with gr.Tab("💾 Profiles"):
                    gr.Markdown(PROFILE_HELP)
                    with gr.Row():
                        profile_picker = gr.Dropdown(
                            label="Saved profiles",
                            choices=[NO_PROFILE, *list_profiles()],
                            value=NO_PROFILE,
                            scale=3,
                        )
                        profile_load = gr.Button("📂 load", variant="secondary", scale=1, elem_id="profile-load")
                        profile_refresh = gr.Button("🔄", variant="secondary", scale=1, min_width=60, elem_id="profile-refresh")
                    with gr.Row():
                        profile_name = gr.Textbox(
                            label="Save as", placeholder="e.g. neon street, turbo 8 steps", max_lines=1, scale=3
                        )
                        profile_save = gr.Button("💾 save", variant="primary", scale=1, elem_id="profile-save")
                    profile_delete = gr.Button("🗑️ delete the selected profile", variant="secondary",
                                               elem_id="profile-delete")
                    profile_status = gr.Markdown("")
                    gr.Markdown("---")
                    gr.Markdown("**Take it with you (.json file)**")
                    save = gr.Button("⬇️ export the current settings to .json", size="sm")
                    settings_download = gr.File(label="Your settings", visible=False, interactive=False)
                    settings_upload = gr.File(
                        label="📤 import a settings .json", file_types=[".json"], type="filepath"
                    )

            with gr.Row():
                estimate = gr.Markdown(elem_classes="gpu-estimate")
                estimate_btn = gr.Button("🔄", variant="secondary", scale=1, min_width=60,
                                       elem_id="estimate-btn")
            run = gr.Button("🚀  Generate", variant="primary", elem_id="run-btn")

        # ---------------- output ----------------
        with gr.Column(scale=6):
            with gr.Group(elem_classes="panel"):
                result = gr.Video(label="🎞️ Video + soundtrack", height=560)
            # An output, so it can be revealed only for a request that asked for a rewrite.
            with gr.Accordion("Upsampled prompt", open=False, visible=False) as upsampled_panel:
                upsampled = gr.Textbox(show_label=False, lines=8, interactive=False)
            with gr.Accordion("🎞️ Stitch clips into one video (no GPU used)", open=False):
                gr.Markdown(
                    "Add clips to a queue and they are joined into a single file, **soundtrack "
                    "included**. Three clips make one long video for no extra GPU time. Everything "
                    "is scaled to the first clip's frame; a clip without audio gets silence rather "
                    "than breaking the join."
                )
                merge_name = gr.Textbox(label="File name (optional)", placeholder="my_scene",
                                        max_lines=1)
                auto_merge = gr.Checkbox(
                    value=False, label="⚡ Add every new generation automatically",
                    info="Generate, generate, generate — the joined file grows on its own.",
                )
                with gr.Row():
                    merge_add_btn = gr.Button("➕ Add the current video", variant="secondary")
                    merge_clear_btn = gr.Button("🧹 Clear the queue", variant="secondary")
                merge_status = gr.Markdown("Queue: empty.")
                merged_video = gr.Video(label="🎬 Stitched result", height=360)
                merged_file = gr.File(label="⬇️ Download the stitched video")
                merge_queue = gr.State([])

            with gr.Accordion("💡 Tips", open=False):
                gr.Markdown(
                    "- the **GPU cost** line above the Generate button is the same figure the Space reserves, so it "
                    "turns red before a request is refused.\n"
                    "- a Turbo lora at 4–8 steps costs roughly a fifth of the default 28 steps.\n"
                    "- a shorter duration and a *fast* canvas are the two biggest savings.\n"
                    "- save a set-up you like as a profile — it comes back from the dropdown next visit."
                )

    open_slots = gr.State(OPEN_IMAGE_SLOTS)

    def reveal_image_slot(open_count):
        open_count = min(open_count + 1, MAX_IMAGE_SLOTS)
        return [
            open_count,
            *[gr.update(visible=index < open_count) for index in range(MAX_IMAGE_SLOTS)],
            gr.update(visible=open_count < MAX_IMAGE_SLOTS),
        ]

    add_image.click(reveal_image_slot, open_slots, [open_slots, *images, add_image], api_name=False)

    for control in (audio, video, match):
        control.change(
            duration_controls, [audio, video, match], [match, duration], show_progress="hidden", api_name=False
        )

    # `reference, strength, reference, strength, ...`, which is how `generate` unpacks them.
    lora_inputs = [field for pair in zip(lora_references, lora_scales) for field in pair]
    lora_upload.upload(_fill_lora_slots, [lora_upload, *lora_references], lora_references, api_name=False)
    lora_preset_add.click(
        _add_preset_lora,
        [lora_preset, *lora_references, *lora_scales],
        [*lora_references, *lora_scales, steps],
        api_name=False,
    )

    # Same order as `SETTINGS_KEYS`.
    settings_fields = [
        prompt, upsample, canvas, match, duration, steps, seed,
        *lora_references, *lora_scales, randomize_seed,
    ]
    save.click(save_settings, settings_fields, settings_download, api_name=False)
    settings_upload.upload(load_settings, settings_upload, settings_fields, api_name=False)

    # Same order as `generate`'s signature: the five leading columns first, then the remaining image slots, then the
    # LoRA fields the `*lora_fields` tail collects.
    request = [
        prompt, images[0], audio, video, canvas, *images[1:], match, duration, steps, seed, upsample, *lora_inputs
    ]

    # A new seed is drawn before the request when the box is ticked, so the number in the box is always the one the
    # clip was made with. `/generate` itself keeps taking the seed it is handed, so the API is unchanged.
    run.click(roll_seed, [randomize_seed, seed], seed, show_progress="hidden", api_name=False).then(
        generate, request, [result, upsampled, upsampled_panel], api_name="generate"
    ).then(
        # CPU-side, so it costs nothing from the GPU allowance.
        auto_queue, [result, auto_merge, merge_queue, merge_name],
        [merged_video, merged_file, merge_queue, merge_status],
        show_progress="hidden", api_name=False,
    )

    merge_add_btn.click(
        add_to_queue, [result, merge_queue, merge_name],
        [merged_video, merged_file, merge_queue, merge_status], api_name=False,
    )
    merge_clear_btn.click(
        clear_queue, None, [merged_video, merged_file, merge_queue, merge_status], api_name=False,
    )
    seed_dice.click(lambda: random.randint(0, MAX_SEED), None, seed, show_progress="hidden", api_name=False)

    # ------------------------------------------------------------------------------------------------------------
    # Named profiles
    # ------------------------------------------------------------------------------------------------------------
    profile_save.click(
        save_profile, [profile_name, *settings_fields], [profile_picker, profile_status], api_name=False
    )
    profile_load.click(
        load_profile, profile_picker, [*settings_fields, profile_name, profile_status], api_name=False
    )
    profile_delete.click(delete_profile, profile_picker, [profile_picker, profile_status], api_name=False)
    profile_refresh.click(refresh_profiles, profile_picker, profile_picker)
    # Repopulate on every page open, so profiles saved elsewhere show up without a restart.
    demo.load(refresh_profiles, profile_picker, profile_picker)

    # ------------------------------------------------------------------------------------------------------------
    # Quick tags, the Turbo blurb, and the live GPU cost
    # ------------------------------------------------------------------------------------------------------------
    def append_chip(text, chip):
        base = (text or "").strip().rstrip(",")
        if chip.lower() in base.lower():
            return base
        return f"{base}, {chip}" if base else chip

    for button, chip in zip(chip_buttons_a + chip_buttons_b, CHIPS):
        button.click(
            (lambda value: (lambda text: append_chip(text, value)))(chip),
            prompt,
            prompt,
            show_progress="hidden",
            api_name=False,
        )

    lora_preset.change(
        lambda name: LORA_PRESETS.get(name, ("", 0, "", 1.0))[2],
        lora_preset,
        turbo_blurb,
        show_progress="hidden",
        api_name=False,
    )

    estimate_inputs = [canvas, duration, steps, match, audio, video, *images, *lora_references]
    # Recalculated on demand only. An automatic readout means an event on every slider
    # move, and a browser that never gets to rest.
    estimate_btn.click(gpu_estimate, estimate_inputs, estimate, show_progress="hidden")
    demo.load(gpu_estimate, estimate_inputs, estimate, api_name=False)


if __name__ == "__main__":
    demo.launch(theme=THEME, css=CSS, show_error=True, ssr_mode=False)