app.py

"""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 gc
import os
import random
import re
import subprocess
import tempfile
import time
import traceback
from functools import cache

# Before torch exists: 72 GiB of resident weights leave the rest of the card in pieces, and a request that needs one
# more large contiguous block then fails on free memory it cannot use in one piece. Expandable segments let the
# allocator grow a block instead of hunting for one, which is the single biggest difference on this Space.
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

# 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

import lora_library

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"))
# An attached adapter adds its own layers and their activations to the same card, so the ceiling above is not the
# ceiling any more. Refusing a request that is over the reduced one is a sentence on screen; letting it through is a
# dead worker and a bare "runtime error".
LORA_SEQUENCE_FRACTION = float(os.environ.get("H3_LORA_SEQUENCE_FRACTION", "0.75"))


def sequence_ceiling(loras=()) -> int:
    return int(MAX_SEQUENCE * LORA_SEQUENCE_FRACTION) if loras else MAX_SEQUENCE
# 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
    )
    ceiling = sequence_ceiling(loras)
    if sequence <= ceiling and total <= MAX_GPU_DURATION:
        return

    seconds = num_frames / FPS
    if sequence > ceiling:
        room = " A lora is attached, which takes part of the card for itself." if loras else ""
        raise gr.Error(
            f"This request is too large for the card: {sequence} rows against a ceiling of {ceiling} "
            f"({width}x{height}, {seconds:.1f} s, {len(references)} references).{room} "
            "Lower the duration, pick a smaller canvas (a 1:1 one is the smallest), 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 _civitai_alternatives(url: str):
    """The same download through each CivitAI domain in turn. They mirror one another, and one of
    them will sometimes answer a datacentre address with a Cloudflare challenge while another does
    not - which is what a Space is, so this is not a rare case."""
    yield url
    host = (urlparse(url).hostname or "").lower()
    if any(host.endswith(known) for known in CIVITAI_HOSTS):
        for other in CIVITAI_HOSTS:
            if not host.endswith(other):
                yield url.replace(host, other, 1)


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

    response, last = None, None
    for candidate in _civitai_alternatives(url):
        candidate_host = (urlparse(candidate).hostname or "").lower()
        attempt_url = candidate
        if token and any(candidate_host.endswith(k) for k in CIVITAI_HOSTS) and "token=" not in candidate:
            attempt_url = candidate + ("&" if "?" in candidate else "?") + f"token={token}"
        try:
            attempt = requests.get(attempt_url, stream=True, timeout=120,
                                   headers={"User-Agent": "Mozilla/5.0"})
            attempt.raise_for_status()
        except Exception as failure:  # noqa: BLE001
            # Not quoting the exception: requests puts the URL in its message, and the URL carries
            # the token.
            last = f"{type(failure).__name__} from {candidate_host}"
            print(f"[lora] {last}; trying another CivitAI domain")
            continue
        if "text/html" in (attempt.headers.get("content-type") or "").lower():
            last = f"a web page instead of a file from {candidate_host}"
            print(f"[lora] {last}; trying another CivitAI domain")
            continue
        response = attempt
        break

    if response is None:
        raise gr.Error(
            f"Could not fetch that link ({last}). CivitAI sometimes puts a Cloudflare challenge in "
            "front of a datacentre address, and a gated model needs CIVITAI_TOKEN under Settings -> "
            "Variables and secrets. Downloading the file yourself and uploading it always works."
        )

    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


# ------------------------------------------------------------------------------------------------------------------
# LoKr conversion
# ------------------------------------------------------------------------------------------------------------------
# LyCORIS LoKr stores a layer as the Kronecker product of two small factors (`lokr_w1` and `lokr_w2`, each possibly
# itself factored into `_a @ _b`), which PEFT cannot load at all. The product is reconstructed and re-expressed as an
# ordinary low-rank pair, exactly: an SVD of a Kronecker product is the outer product of the factors' SVDs, so the
# largest `rank` singular values can be picked without ever building the full matrix - which for H3 would be 7168 by
# 7168 per layer. What survives is reported as a percentage; a low number means the rank was too small to hold the
# adapter, not that anything went wrong.

LOKR_RANK = int(os.environ.get("H3_LOKR_RANK", "32"))
_LOKR_SUFFIXES = ("lokr_w1_a", "lokr_w1_b", "lokr_w2_a", "lokr_w2_b", "lokr_t2", "lokr_w1", "lokr_w2", "alpha")
_LOKR_PREFIXES = ("lycoris_", "lycoris.")


def _is_lokr_lora(state_dict) -> bool:
    return any(".lokr_w2" in key or ".lokr_w1" in key for key in state_dict)


def _lokr_module_name(module: str):
    """`lycoris_blocks_0_attn_qkv_proj` -> `blocks.0.attn.qkv_proj`, which the kohya pass then renames."""
    name = module
    for prefix in _LOKR_PREFIXES:
        if name.startswith(prefix):
            name = name[len(prefix):]
    name = _strip_lora_prefixes(name)
    if "blocks." in name:
        return name
    dotted = _dotted_module(name)
    return dotted


def _kron_low_rank(w1, w2, rank: int, scale: float):
    """`(A, B, kept energy)` such that `B @ A` approximates `scale * kron(w1, w2)`.

    `svd(kron(w1, w2))` has singular values `outer(s1, s2)` and vectors `kron(u1_i, u2_j)`, so the truncation is a
    choice among those products rather than a decomposition of the big matrix.
    """
    import torch

    u1, s1, v1 = torch.linalg.svd(w1.float(), full_matrices=False)
    u2, s2, v2 = torch.linalg.svd(w2.float(), full_matrices=False)

    products = torch.outer(s1, s2) * scale
    flat = products.reshape(-1)
    keep = int(min(rank, flat.numel()))
    order = torch.argsort(flat, descending=True)[:keep]

    rows = torch.div(order, s2.numel(), rounding_mode="floor")
    cols = order % s2.numel()
    sigma = torch.sqrt(torch.clamp(flat[order], min=0.0))

    lora_b = torch.stack(
        [torch.outer(u1[:, i], u2[:, j]).reshape(-1) * s for i, j, s in zip(rows, cols, sigma)], dim=1
    )
    lora_a = torch.stack(
        [torch.outer(v1[i, :], v2[j, :]).reshape(-1) * s for i, j, s in zip(rows, cols, sigma)], dim=0
    )

    total = float(flat.sum())
    energy = float(flat[order].sum() / total) if total > 0 else 1.0
    return lora_a, lora_b, energy


def _convert_lokr_lora(state_dict, rank: int = None) -> dict:
    """LoKr -> the `lora_down` / `lora_up` pairs the kohya pass understands. Returns `{}` when nothing matched."""
    import torch

    rank = int(rank or LOKR_RANK)
    modules = {}
    for key in state_dict:
        for suffix in _LOKR_SUFFIXES:
            if key.endswith("." + suffix):
                modules.setdefault(key[: -(len(suffix) + 1)], {})[suffix] = key
                break

    converted, worst, done, skipped = {}, 1.0, 0, 0
    for module, roles in sorted(modules.items()):
        if "lokr_w2" not in roles and "lokr_w2_b" not in roles:
            continue

        name = _lokr_module_name(module)
        if not name or "lokr_t2" in roles:
            skipped += 1
            continue

        inner = None
        if "lokr_w1" in roles:
            w1 = state_dict[roles["lokr_w1"]].float()
        else:
            w1_b = state_dict[roles["lokr_w1_b"]].float()
            inner = w1_b.shape[0]
            w1 = state_dict[roles["lokr_w1_a"]].float() @ w1_b

        if "lokr_w2" in roles:
            w2 = state_dict[roles["lokr_w2"]].float()
        else:
            w2_b = state_dict[roles["lokr_w2_b"]].float()
            inner = w2_b.shape[0]
            w2 = state_dict[roles["lokr_w2_a"]].float() @ w2_b

        if w1.ndim != 2 or w2.ndim != 2:
            skipped += 1
            continue

        # LyCORIS scales by alpha / inner-dim, the same convention kohya uses for its rank.
        scale = 1.0
        if inner and "alpha" in roles:
            try:
                scale = float(state_dict[roles["alpha"]].reshape(-1)[0]) / inner
            except Exception:  # noqa: BLE001
                scale = 1.0

        keep = int(min(rank, w1.shape[0] * min(w2.shape)))
        lora_a, lora_b, energy = _kron_low_rank(w1, w2, keep, scale)
        worst = min(worst, energy)
        done += 1

        converted[f"{name}.lora_down.weight"] = lora_a.to(torch.float32)
        converted[f"{name}.lora_up.weight"] = lora_b.to(torch.float32)

    if converted:
        print(f"[lora] LoKr conversion: {done} layers at rank {rank}, weakest layer keeps "
              f"{worst * 100:.0f}% of its strength, skipped {skipped}")
        if worst < 0.6:
            print("[lora] a weak layer means the rank is too small for this adapter - raise H3_LOKR_RANK")
    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.

    A kohya / CivitAI file can carry the same dotted block names, and its fused QKV is interleaved per head rather
    than stored as three thirds, so the two must not be confused. `lora_down` / `lora_up` / `alpha` are kohya's own
    naming and settle it: ComfyUI's Turbo files are `lora_A` / `lora_B` with no alpha.
    """
    keys = list(state_dict)
    if not any(key.startswith(("blocks.", "token_refiner.blocks.", "final_layer.")) for key in keys):
        return False
    kohya_shaped = any(
        ".lora_down" in key or ".lora_up" in key or key.endswith(".alpha") or key.endswith(".lora_alpha")
        for key in keys
    )
    return not kohya_shaped


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


_CONTAINER_PREFIXES = (
    "model.diffusion_model.", "diffusion_model.", "base_model.model.", "base_model.",
    "transformer_ref.", "transformer.",
)

# The feed-forward is the one module whose name genuinely differs between the trainers *and* between diffusers
# versions of the H3 block. Whichever spelling a file arrives in, it is remapped to the one the loaded transformer
# actually has, so a mismatch can no longer surface as a load-time RuntimeError.
_FF_ALIASES = (
    (".ff.net.0.proj", ".ff.fc1"),
    (".ff.net.2", ".ff.fc2"),
    (".ff.fc1", ".ff.net.0.proj"),
    (".ff.fc2", ".ff.net.2"),
    (".ff.gate_proj", ".ff.fc1"),
    (".ff.down_proj", ".ff.fc2"),
    (".attn.to_out.0", ".attn.out_proj"),
    (".attn.out_proj", ".attn.to_out.0"),
)


def _strip_container_prefix(state_dict) -> dict:
    """Drop a wrapper prefix (`diffusion_model.`, `transformer.`, ...) so format detection sees the real names."""
    for prefix in _CONTAINER_PREFIXES:
        if all(key.startswith(prefix) for key in state_dict):
            return {key[len(prefix):]: value for key, value in state_dict.items()}
    return dict(state_dict)


def _linear_module_names(transformer) -> set:
    """Every module name on the transformer. Anything an adapter can target is in here."""
    try:
        return {name for name, _ in transformer.named_modules() if name}
    except Exception:  # noqa: BLE001
        return set()


def _fit_to_transformer(transformer, state_dict) -> dict:
    """Point every converted key at a module the transformer really has, and drop the ones it does not.

    `load_lora_adapter` raises on a key whose module is missing, which is what a Turbo adapter written against a
    different spelling of the feed-forward looks like from the outside: a runtime error with no useful message.
    """
    targets = _linear_module_names(transformer)
    if not targets:
        return state_dict

    fitted, renamed, dropped = {}, 0, 0
    for key, value in state_dict.items():
        module = key
        for suffix in (".lora_A.weight", ".lora_B.weight", ".lora_A", ".lora_B", ".alpha", ".weight"):
            if module.endswith(suffix):
                module = module[: -len(suffix)]
                break
        leaf = key[len(module):]

        if module in targets:
            fitted[key] = value
            continue

        moved = None
        for old, new in _FF_ALIASES:
            if old in module and module.replace(old, new) in targets:
                moved = module.replace(old, new)
                break
        if moved is None:
            dropped += 1
            continue
        fitted[f"{moved}{leaf}"] = value
        renamed += 1

    if renamed or dropped:
        print(f"[lora] fitted to the transformer: {renamed} keys renamed, {dropped} dropped, {len(fitted)} kept")
    return fitted or state_dict


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)
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()

    base_dtype = next(
        (param.dtype for key, param in transformer.named_parameters() if ".lora_" not in key), torch.bfloat16
    )

    names, scales = [], []
    for index, (path, scale) in enumerate(loras):
        state_dict = _strip_container_prefix(_load_lora_state_dict(path))
        if _is_lokr_lora(state_dict):
            # LoKr first: it produces kohya-shaped keys, which the pass below then renames to diffusers'.
            state_dict = _convert_lokr_lora(state_dict) or state_dict
        if not any(marker in key for key in state_dict for marker in _DIFFUSERS_MARKERS):
            # ComfyUI first. Its fused QKV is three plain thirds while kohya's is interleaved per attention head,
            # and the kohya pass recognises ComfyUI's dotted block names too - so running it first splits a Turbo
            # adapter the wrong way and renames its feed-forward to a module that is not there.
            if _is_comfyui_lora(state_dict):
                state_dict = _convert_comfyui_lora(state_dict)
            else:
                state_dict = _convert_kohya_lora(state_dict) or state_dict
        state_dict = _fit_to_transformer(transformer, state_dict)
        # Cast on the host, before anything crosses onto the card. A float32 adapter injected as-is is put on the
        # card at float32 and cast afterwards, so for a moment both copies are resident - on top of 72 GiB of
        # weights that is exactly the block the allocator cannot find.
        state_dict = {
            key: (value.to(base_dtype) if hasattr(value, "is_floating_point") and value.is_floating_point()
                  else value)
            for key, value in state_dict.items()
        }
        name = f"lora{index}"
        try:
            transformer.load_lora_adapter(state_dict, adapter_name=name, prefix=_lora_prefix(state_dict))
        except Exception as error:  # noqa: BLE001
            traceback.print_exc()
            raise gr.Error(
                f"`{os.path.basename(path)}` could not be attached: `{type(error).__name__}: {error}`. "
                "It is almost always an adapter for another partition or another base model - a `transformer/` "
                "adapter does not fit the `transformer_ref/` half this Space runs."
            ) from error
        del state_dict
        gc.collect()
        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)
    # The load and the cast both leave freed blocks behind; handing them back before the denoise starts is what
    # keeps the first large allocation of the run from failing on a card that has the memory but not in one piece.
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
    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")

    attached = apply_loras(PIPE.transformer_ref, loras or ())

    # cuDNN's fused attention takes q/k/v straight from the projections. With PEFT in front of them the tensors
    # arrive with a different layout, and the kernel refuses shapes it accepted before — which is why the same
    # request can fail on a 16:9 canvas and go through on a square one. SDPA has no such restriction, so an
    # adapter-carrying request runs on `native` and everything else keeps the faster backend.
    def _backend(name):
        try:
            PIPE.transformer_ref.set_attention_backend(name)
        except Exception:  # noqa: BLE001
            pass

    _backend("native" if attached else ATTENTION)

    def _run():
        return 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)),
        )

    try:
        state = _run()
    except RuntimeError as error:
        message = str(error).lower()
        kernel = "no available kernel" in message or "cudnn" in message
        # `NVML_SUCCESS == r INTERNAL ASSERT FAILED` out of `CUDACachingAllocator.cpp` is not a bug to report to
        # PyTorch: it is the allocator failing to grow, i.e. out of memory, with its own error path falling over on
        # the way to saying so. Same treatment as a plain OOM.
        starved = (
            "out of memory" in message
            or "nvml" in message
            or "cudacachingallocator" in message
            or "cuda error" in message
        )
        if not (kernel or starved):
            raise

        gc.collect()
        try:
            torch.cuda.empty_cache()
            torch.cuda.synchronize()
        except Exception:  # noqa: BLE001
            pass
        _backend("native")
        print(f"[ref2va] first attempt failed ({type(error).__name__}); retrying once on a cleared card", flush=True)

        try:
            state = _run()
        except RuntimeError as second:
            second_message = str(second).lower()
            if not (
                "out of memory" in second_message
                or "nvml" in second_message
                or "cudacachingallocator" in second_message
            ):
                raise
            traceback.print_exc()
            raise gr.Error(
                f"The card ran out of memory at {width}x{height}, {num_frames / FPS:.1f} s, "
                f"{len(references)} reference(s)"
                f"{f' and {len(loras or ())} lora' if loras else ''}. "
                "A lora takes its share of the same card, so a setting that worked without one can be too much "
                "with it. Pick a smaller canvas — a 1:1 one is the smallest — shorten the duration, or drop a "
                "reference, and it goes through."
            ) from second

    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

    ceiling = sequence_ceiling(lora_fields)
    if sequence > ceiling:
        head = (
            f"🚫 **Too large for the card:** {sequence} rows against a ceiling of {ceiling}. "
            "Lower the duration, pick a smaller canvas (1:1 is the smallest), 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>"


CIVITAI_DOWNLOAD_RE = re.compile(r"/api/download/models/(\d+)")


@cache
def civitai_details(version_id: str, file_id: str = ""):
    """`(label, trigger words)` for a CivitAI download link, from its public model-versions endpoint.

    The number in a download URL is the *model version* id, so one lookup gives the model's title, the version name,
    the file behind `fileId`, and the words the adapter was trained on. Public, cached, and never fatal: a link that
    cannot be identified simply keeps showing its number.
    """
    import requests

    headers = {"User-Agent": "Mozilla/5.0"}
    token = os.environ.get("CIVITAI_TOKEN", "").strip()
    if token:
        headers["Authorization"] = f"Bearer {token}"

    try:
        response = requests.get(
            f"https://civitai.com/api/v1/model-versions/{version_id}", headers=headers, timeout=20
        )
        response.raise_for_status()
        data = response.json()
    except Exception as error:  # noqa: BLE001
        return f"CivitAI {version_id} (name unavailable: {type(error).__name__})", []

    model_name = (data.get("model") or {}).get("name") or f"model {data.get('modelId', '?')}"
    version_name = data.get("name") or ""
    label = f"{model_name} · {version_name}".strip(" ·")

    if file_id:
        for entry in data.get("files") or []:
            if str(entry.get("id")) == str(file_id):
                label += f" · {entry.get('name', '')}"
                break

    return label, [word for word in (data.get("trainedWords") or []) if word]


def describe_lora(reference: str):
    """A readable line for whatever is in a slot: a CivitAI link becomes its real title."""
    reference = (reference or "").strip()
    if not reference:
        return "", []

    match = CIVITAI_DOWNLOAD_RE.search(reference)
    if match:
        file_id = ""
        if "fileId=" in reference:
            file_id = reference.split("fileId=", 1)[1].split("&")[0]
        return civitai_details(match.group(1), file_id)

    if reference.startswith(("http://", "https://")):
        return os.path.basename(reference.split("?")[0]) or reference, []
    return reference, []


def identify_loras(*references):
    """Name every filled slot. CivitAI links are looked up; everything else is shown as typed."""
    lines = []
    for index, reference in enumerate(references, start=1):
        reference = (reference or "").strip()
        if not reference:
            continue
        label, words = describe_lora(reference)
        line = f"**{index}.** {label}"
        if words:
            line += f"  \n<sub>trigger words: {', '.join(words[:8])}</sub>"
        lines.append(line)
    if not lines:
        return "Nothing in the slots yet."
    return "  \n".join(lines)


# ------------------------------------------------------------------------------------------------------------------
# Continuing a scene, and searching CivitAI
# ------------------------------------------------------------------------------------------------------------------


def _frames_rgb(video_path, first_only: bool = False, keep: int = 4):
    """Decoded frames as RGB arrays. `first_only` stops after the opening frame; otherwise
    the last `keep` frames are returned, decoded forward rather than seeked to."""
    import av

    frames = []
    with av.open(str(video_path)) as container:
        stream = container.streams.video[0]
        stream.thread_type = "AUTO"
        for frame in container.decode(stream):
            frames.append(frame.to_ndarray(format="rgb24"))
            if first_only:
                break
            if len(frames) > keep:
                frames.pop(0)
    return frames


def last_frame_of(video_path) -> str:
    """Write the final frame of a clip to a PNG and return its path, so it can be dropped straight into an image
    slot. The true last frame: stepping back a few frames "for safety" is what put the join out by roughly one
    frame, because the next clip then began from a moment the previous clip had already played past."""
    if not video_path or not os.path.exists(str(video_path)):
        raise gr.Error("Generate a video first - there is nothing to continue from.")

    from PIL import Image

    frames = _frames_rgb(video_path)
    if not frames:
        raise gr.Error("That video has no readable frames.")

    # The genuine end of the clip, unless it really is blank - then step back to the last one
    # that carries an image.
    chosen = frames[-1]
    for candidate in reversed(frames):
        if float(candidate.mean()) > 2.0:
            chosen = candidate
            break

    directory = os.path.join(tempfile.gettempdir(), "continuations")
    os.makedirs(directory, exist_ok=True)
    path = os.path.join(directory, f"frame_{int(time.time() * 1000)}.png")
    Image.fromarray(chosen).save(path)
    return path


def stage_extension(video_path, queue, name_hint):
    """Before a continuation runs: park the finished clip in the merge queue and hand its last frame back as the new
    first reference."""
    frame = last_frame_of(video_path)
    merged, merged_file, queue, status = add_to_queue(video_path, queue, name_hint)
    return frame, None, merged, merged_file, queue, status


# Which host answers the search. `.red` is a mirror of the same API and carries entries the main
# domain hides, so it is asked first and `.com` is the fallback. Set CIVITAI_API_HOST to pin one.
# Note that adult entries are returned to an authenticated caller only, whichever host answers:
# without CIVITAI_TOKEN the X / XXX levels are simply absent from the results.
CIVITAI_API_HOSTS = [h for h in (os.environ.get("CIVITAI_API_HOST", "").strip(),
                                 "civitai.red", "civitai.com") if h]


H3_BASE_MODELS = ["MiniMax H3", "(any base model)"]


def civitai_search(query, base_model, want_nsfw, limit=20):
    """Search CivitAI for lora. `/api/v1/models` embeds each model's versions, files and trigger words, so one call
    gives everything a slot needs. Returns `(readable list, dropdown update, {label: url})`."""
    import requests

    query = (query or "").strip()
    if not query:
        return "Type something to search for.", gr.update(choices=[], value=None), {}

    params = {"query": query, "types": "LORA", "limit": int(limit), "sort": "Most Downloaded"}
    if base_model and base_model != "(any base model)":
        params["baseModels"] = base_model
    if want_nsfw:
        params["nsfw"] = "true"

    headers = {"User-Agent": "Mozilla/5.0"}
    token = os.environ.get("CIVITAI_TOKEN", "").strip()
    if token:
        # Both forms on purpose. The header is the documented one, but search honours the query
        # parameter more reliably - and without an authenticated call the adult browsing levels are
        # simply missing from the results, whichever host answers.
        headers["Authorization"] = f"Bearer {token}"
        params["token"] = token

    items, error, answered = [], None, False
    host_used = CIVITAI_API_HOSTS[0]
    for host in CIVITAI_API_HOSTS:
        try:
            response = requests.get(f"https://{host}/api/v1/models", params=params,
                                    headers=headers, timeout=30)
            response.raise_for_status()
            items = response.json().get("items") or []
            host_used, answered = host, True
            if items:
                break
        except Exception as failure:  # noqa: BLE001
            error = failure

    if not items and error is not None and not answered:  # noqa: BLE001
        # Deliberately not quoting the exception: `requests` puts the full URL in its message, and the
        # URL carries the token as a query parameter, so echoing it would print the key on screen.
        status = getattr(getattr(error, "response", None), "status_code", None)
        reason = f"HTTP {status}" if status else type(error).__name__
        tried = ", ".join(CIVITAI_API_HOSTS)
        return (f"Search failed ({reason}). Tried: {tried}. CivitAI answers 503 when it is rate-limiting "
                f"or briefly down - wait a moment and press Search again.",
                gr.update(choices=[], value=None), {})

    if not items:
        return ("Nothing found. Try fewer words, or set the base model to *(any base model)*.",
                gr.update(choices=[], value=None), {})

    mapping, lines, choices = {}, [], []
    for item in items:
        model_name = item.get("name") or "?"
        creator = (item.get("creator") or {}).get("username") or "?"
        downloads = (item.get("stats") or {}).get("downloadCount") or 0

        for version in (item.get("modelVersions") or [])[:3]:
            version_name = version.get("name") or ""
            # What the adapter was trained against. Worth showing on every row: with the filter set to
            # "(any base model)" the results mix families, and a lora for another one downloads happily
            # and then does nothing useful.
            base = version.get("baseModel") or "base model unknown"
            words = [w for w in (version.get("trainedWords") or []) if w]
            for entry in (version.get("files") or []):
                name = entry.get("name") or ""
                if not name.lower().endswith(".safetensors"):
                    continue
                url = f"https://{host_used}/api/download/models/{version.get('id')}?fileId={entry.get('id')}"
                size = float(entry.get("sizeKB") or 0) / 1024
                label = f"{len(choices) + 1}. [{base}] {model_name} · {version_name} · {name}"[:160]
                mapping[label] = url
                choices.append(label)
                line = (f"**{label.split('. ', 1)[0]}.** **{model_name}** · {version_name} · {name}  \n"
                        f"<sub>**{base}** · {size:.0f} MB · {downloads} downloads · by {creator}")
                if words:
                    line += f" · triggers: {', '.join(words[:5])}"
                lines.append(line + "</sub>")

    if not choices:
        return "Found models, but none with a `.safetensors` file.", gr.update(choices=[], value=None), {}

    # Wrapped so the list scrolls in place instead of pushing the page down.
    body = "<div class=\"search-results\">\n\n" + "  \n".join(lines[:40]) + "\n\n</div>"
    return body, gr.update(choices=choices, value=choices[0]), mapping


def put_in_slot(label, mapping, slot):
    """Write the chosen result's download link into one lora slot, leaving the others alone."""
    updates = [gr.update() for _ in range(LORA_SLOTS)]
    url = (mapping or {}).get(label)
    if not url:
        return [*updates, "Search and pick a file first."]

    index = int(str(slot).split()[-1]) - 1
    updates[index] = gr.update(value=url)
    return [*updates, f"**Put into {slot}** — {label}"]


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

# ------------------------------------------------------------------------------------------------------------------
# Structured prompt builder
# ------------------------------------------------------------------------------------------------------------------
# H3 was trained on the output of H3-Context-IR, a preprocessor that rewrites a plain request into labelled sections,
# and MiniMax's own model card calls that structure "critical to the quality of the final output". Nothing in this
# pipeline adds it: the string reaches the transformer as typed. So the builder writes the sections instead - the
# shot description, the soundscape and the music, in the order and under the names the model was trained to read.
#
# Two details from the official guide are worth knowing, because getting them wrong looks like a model fault:
# dialogue must be verbatim inside <d> tags or the mouth moves with no words in it, and reference tags have to
# appear in the order the inputs were connected.

IR_SHOT_TYPES = {
    "(none)": "",
    "live-action, cinematic": "Live-action, cinematic",
    "live-action, documentary": "Live-action, documentary, handheld",
    "studio portrait": "Live-action, studio portrait lighting",
    "anime": "2D anime, crisp lineart",
    "3D animation": "3D animation, stylised",
}

IR_CAMERA = {
    "(none)": "",
    "static": "The camera holds a static frame",
    "slow push in": "The camera pushes in with small amplitude at slow speed",
    "slow pull back": "The camera pulls back with small amplitude at slow speed",
    "truck right": "The camera trucks right with small amplitude at slow speed",
    "truck left": "The camera trucks left with small amplitude at slow speed",
    "orbit": "The camera orbits the subject with medium amplitude at slow speed",
    "handheld follow": "The camera follows handheld with small amplitude at moderate speed",
    "tilt up": "The camera tilts up with small amplitude at slow speed",
    "crane up and back": "The camera cranes up and back with large amplitude at slow speed",
}

IR_SOUNDSCAPE = {
    "(none)": "",
    "quiet room": "A quiet room tone with small incidental sounds - fabric, footsteps, a distant door.",
    "rain and traffic": "Rain ticks against glass over the low hum of distant traffic.",
    "outdoors, wind": "Wind moves through the scene, carrying faint birdsong and rustling leaves.",
    "city street": "City ambience: passing cars, footsteps on pavement, indistinct voices further off.",
    "interior, machinery": "A steady mechanical hum underneath, with occasional metallic ticks.",
    "crowd": "A crowd murmurs at a middle distance, individual voices indistinct.",
}

IR_MUSIC = {
    "(none)": "",
    "no music": "None.",
    "slow strings": "Sustained cello notes at a slow tempo with widely spaced piano tones.",
    "warm piano": "A warm solo piano at a slow tempo, sparse and unhurried.",
    "tense low drone": "A low synth drone with a slow rising tension.",
    "upbeat electronic": "An upbeat electronic pulse at a moderate tempo.",
}


def build_ir_prompt(description, shot_type, camera, soundscape, music, speaker, dialogue,
                    reference_count):
    """Write the labelled sections H3 was trained on, around what the user typed."""
    body = (description or "").strip().rstrip(".")
    if not body:
        raise gr.Error("Describe the shot first - the builder writes the structure around it.")

    lines = []

    # The reference line comes first, and names the pictures in connection order, which is what the
    # model expects to match against its inputs.
    count = int(reference_count or 0)
    if count > 0:
        tags = ", ".join(f"<Picture {index}>" for index in range(1, count + 1))
        lines.append(
            f"For the target video, at 0.00 seconds into the target video, {tags} "
            f"(from [Shot 1]) {'is' if count == 1 else 'are'} fully referenced."
        )
        lines.append("")

    shot = " ".join(part for part in (IR_SHOT_TYPES.get(shot_type, ""),) if part)
    described = f"{shot}, {body}." if shot else f"{body}."
    if count > 0:
        described += (" The subject keeps the appearance shown in the reference images, and the "
                      "setting keeps its layout.")
    if IR_CAMERA.get(camera):
        described += f" {IR_CAMERA[camera]}."

    if (dialogue or "").strip():
        who = (speaker or "S1").strip() or "S1"
        spoken = dialogue.strip().strip('"')
        # Delivery and identity go outside the tag; only the language and the verbatim words go in,
        # otherwise the mouth moves correctly with nothing in it.
        described += f" The speaker ({who}) says: <d>[English] {spoken}</d>"

    lines.append(f"integrated_multimodal_description: [Shot 1] {described}")
    lines.append("")
    lines.append(f"overall_soundscape: {IR_SOUNDSCAPE.get(soundscape) or 'Ambient sound suited to the scene.'}")
    lines.append("")
    lines.append(f"non_diegetic_music: {IR_MUSIC.get(music) or 'None.'}")

    return "\n".join(lines)


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 _thumb(array):
    """A tiny, cheap fingerprint of a frame - enough to tell a repeat from a new shot."""
    if array is None:
        return None
    try:
        import numpy as np

        height, width = array.shape[:2]
        rows = np.linspace(0, height - 1, 48).astype(int)
        cols = np.linspace(0, width - 1, 48).astype(int)
        return array[rows][:, cols].astype("float32")
    except Exception:  # noqa: BLE001
        return None


def _same_frame(a, b, tolerance: float = 7.0) -> bool:
    if a is None or b is None or a.shape != b.shape:
        return False
    try:
        import numpy as np

        return float(np.abs(a - b).mean()) < tolerance
    except Exception:  # noqa: BLE001
        return False


def _head_thumb(path):
    try:
        frames = _frames_rgb(path, first_only=True)
        return _thumb(frames[0]) if frames else None
    except Exception:  # noqa: BLE001
        return None


def _tail_thumb(path):
    try:
        frames = _frames_rgb(path)
        return _thumb(frames[-1]) if frames else None
    except Exception:  # noqa: BLE001
        return None


def _normalise(path, width, height, out_path, trim_head: bool = False):
    """One frame size, one frame rate, one audio format. `trim_head` drops the opening frame
    and the matching sliver of audio, so picture and sound stay locked together."""
    video_chain = (f"scale={width}:{height}:force_original_aspect_ratio=decrease,"
                   f"pad={width}:{height}:(ow-iw)/2:(oh-ih)/2,setsar=1,fps={FPS}")
    audio_chain = "aresample=48000,aformat=sample_fmts=fltp:channel_layouts=stereo"
    if trim_head:
        video_chain += ",trim=start_frame=1,setpts=PTS-STARTPTS"
        audio_chain += f",atrim=start={1.0 / FPS:.6f},asetpts=PTS-STARTPTS"

    _, _, has_audio = _probe(path)
    if has_audio:
        filter_str = f"[0:v]{video_chain}[v];[0:a]{audio_chain}[a]"
        command = [_ffmpeg_exe(), "-y", "-i", path,
                   "-filter_complex", filter_str, "-map", "[v]", "-map", "[a]",
                   "-c:v", "libx264", "-preset", "fast", "-crf", "20", "-pix_fmt", "yuv420p",
                   "-c:a", "aac", "-ar", "48000", "-ac", "2", out_path]
    else:
        filter_str = f"[0:v]{video_chain}[v];[1:a]{audio_chain}[a]"
        command = [_ffmpeg_exe(), "-y", "-i", path,
                   "-f", "lavfi", "-i", "anullsrc=r=48000:cl=stereo",
                   "-filter_complex", filter_str, "-map", "[v]", "-map", "[a]", "-shortest",
                   "-c:v", "libx264", "-preset", "fast", "-crf", "20", "-pix_fmt", "yuv420p",
                   "-c:a", "aac", "-ar", "48000", "-ac", "2", 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")

    # The seam. A continued clip opens on the still lifted off the end of the clip before it,
    # so that frame sits in both files and the join stutters by exactly one frame. Compare the
    # two and drop the repeat only when it is really there.
    trims = [False] * len(valid)
    previous_tail = _tail_thumb(valid[0])
    for index in range(1, len(valid)):
        trims[index] = _same_frame(previous_tail, _head_thumb(valid[index]))
        previous_tail = _tail_thumb(valid[index])

    try:
        parts = [
            _normalise(path, width, height, os.path.join(work_dir, f"part{index}.mp4"),
                       trim_head=trims[index])
            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]", "-r", str(FPS),
             "-c:v", "libx264", "-preset", "fast", "-crf", "20", "-pix_fmt", "yuv420p",
             "-c:a", "aac", "-ar", "48000", "-ac", "2",
             "-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 [])
    # Not twice. "Continue the scene" parks the finished clip before it generates, and the previous
    # press already parked that same clip on its way out - so without this the queue grows
    # A, B, B, C, C rather than A, B, C.
    if video_path and os.path.exists(str(video_path)) and (not queue or queue[-1] != 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;
}

.search-results {
    max-height: 300px;
    overflow-y: auto;
    border-radius: 12px;
    padding: 4px 14px;
    background: var(--background-fill-secondary);
    border: 1px solid var(--border-color-primary);
    font-size: 13px;
}

.consent-note {
    border-radius: 12px;
    padding: 8px 12px;
    margin-bottom: 10px;
    background: rgba(239, 68, 68, .08);
    border: 1px solid rgba(239, 68, 68, .30);
    font-size: 12.5px;
    opacity: .9;
}

.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);
}
#ir-btn, #search-btn, #search-put, #extend-btn, #lora-identify, #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; Shared lora library + CivitAI search,
  structured prompt builder, scene continuation, 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. The prompt builder writes the labelled sections H3 was actually trained on, dialogue tags
  and all, and the shared lora library remembers every adapter anyone adds &mdash; link, trigger words and strength
  &mdash; and fills the slots in one press.
  <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>shared lora library</span><span>trigger words in the prompt</span><span>structured prompt builder</span><span>named profiles</span><span>CivitAI search + links</span><span>scene continuation</span><span>kohya + LoKr 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.

A CivitAI link is just a number, so press **🔎 name the links** (or Enter in a slot) and each one is replaced by its
real title, version, file name and trigger words, read from CivitAI's public API.

**LoKr adapters work too.** LyCORIS LoKr stores a layer as the Kronecker product of two small factors, which PEFT
cannot load at all. It is rebuilt into an ordinary low-rank pair at load time — exactly, since the SVD of a Kronecker
product is the outer product of the factors' SVDs, so nothing the size of the full 7168×7168 layer is ever built. The
log reports how much of the weakest layer survived the truncation; if that number is low, raise `H3_LOKR_RANK`
(default 32) under Settings → Variables and secrets.

**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, structured prompts, GPU cost, profiles, 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("🎬 Structured prompt builder (what H3 was trained on)", open=False):
                    gr.Markdown(
                        "H3 was trained on the output of a preprocessor that rewrites a request into labelled "
                        "sections, and MiniMax call that structure *critical to the quality of the final output*. "
                        "Nothing here adds it on its own — your text reaches the model as typed. Describe the shot "
                        "in the prompt box above, set the pieces below, and press **Build**.\n\n"
                        "**Dialogue has to be verbatim.** Speech is generated together with the picture, so naming "
                        "that someone speaks without giving the words produces correct mouth shapes with nothing in "
                        "them. Aim for 350–500 words of description for a full scene."
                    )
                    with gr.Row():
                        ir_shot = gr.Dropdown(list(IR_SHOT_TYPES), value="live-action, cinematic",
                                              label="Shot type")
                        ir_camera = gr.Dropdown(list(IR_CAMERA), value="slow push in", label="Camera move")
                    with gr.Row():
                        ir_sound = gr.Dropdown(list(IR_SOUNDSCAPE), value="(none)", label="Soundscape")
                        ir_music = gr.Dropdown(list(IR_MUSIC), value="no music", label="Music")
                    with gr.Row():
                        ir_speaker = gr.Textbox(value="S1", label="Speaker id", max_lines=1, scale=1)
                        ir_dialogue = gr.Textbox(label="Dialogue, word for word", scale=4,
                                                 placeholder="I get off at the next station.")
                    ir_references = gr.Slider(0, MAX_IMAGE_SLOTS, value=1, step=1,
                                              label="Reference images to name",
                                              info="Named in connection order, as <Picture 1>, <Picture 2> …")
                    ir_button = gr.Button("🎬 Build the structured prompt", variant="secondary",
                                          elem_id="ir-btn")

                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.
                        gr.Markdown(
                            "Upload only photographs of yourself, or of people who have agreed to this. Sexual "
                        "imagery of a real person made without their consent breaks Hugging Face's "
                        "content policy and is not welcome here.",
                            elem_classes="consent-note",
                        )
                        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)
                    with gr.Accordion("🔍 Search CivitAI", open=False):
                        gr.Markdown(
                            "Search CivitAI without leaving the Space, then drop a result straight into a slot. "
                            "Each hit shows its size, downloads and trigger words."
                        )
                        with gr.Row():
                            search_query = gr.Textbox(label="Search", placeholder="e.g. dance, rain, camera move",
                                                      scale=3)
                            search_base = gr.Dropdown(H3_BASE_MODELS, value=H3_BASE_MODELS[0], label="Base model",
                                                      allow_custom_value=True, scale=2)
                        with gr.Row():
                            search_nsfw = gr.Checkbox(value=True, label="Include NSFW results")
                            search_btn = gr.Button("🔍 Search", variant="secondary", elem_id="search-btn")
                        search_pick = gr.Dropdown(choices=[], label="Pick a file")
                        with gr.Row():
                            search_slot = gr.Dropdown([f"lora {i}" for i in range(1, LORA_SLOTS + 1)],
                                                      value="lora 1", label="Into slot", scale=2)
                            search_put_btn = gr.Button("⬇️ Put it in", variant="primary", scale=1,
                                                       elem_id="search-put")
                        # Below the picker on purpose: a wall of results above it would push the
                        # controls off the screen, which is exactly what happened the first time.
                        search_results = gr.Markdown("No search yet.")
                        search_state = gr.State({})

                    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,
                                )
                            )
                    with gr.Row():
                        lora_identify_btn = gr.Button("🔎 name the links", variant="secondary",
                                                      elem_id="lora-identify")
                    lora_names = gr.Markdown("Nothing in the slots yet.", elem_classes="turbo-blurb")
                    lora_upload = gr.File(
                        label="Drop .safetensors here to fill the slots",
                        file_count="multiple",
                        file_types=[".safetensors"],
                        type="filepath",
                    )

                lora_library.library_tab(lora_slots=lora_references, scale_slots=lora_scales,
                                         prompt_box=prompt)

                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)
                extend_btn = gr.Button("➕ Continue the scene (one more clip, joined on)",
                                       variant="secondary", elem_id="extend-btn")
                gr.Markdown(
                    "Takes the last frame of the clip above, makes it the first reference, generates again with "
                    "the settings untouched, and joins the two into one file — soundtrack included. Press it again "
                    "for a third clip. Each press costs one normal generation; the joining is free.",
                    elem_classes="turbo-blurb",
                )
            # 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,
    )

    # Continuing a scene: park the finished clip, hand its last frame to the first image slot, generate from it,
    # then join the new clip on. `request` starts with the prompt and the first image, so the second step reads
    # the still the first step just wrote.
    extend_btn.click(
        fn=stage_extension,
        inputs=[result, merge_queue, merge_name],
        outputs=[images[0], result, merged_video, merged_file, merge_queue, merge_status],
        api_name=False,
    ).then(
        roll_seed, [randomize_seed, seed], seed, show_progress="hidden", api_name=False,
    ).then(
        generate, request, [result, upsampled, upsampled_panel], api_name=False,
    ).then(
        auto_queue, [result, auto_merge, merge_queue, merge_name],
        [merged_video, merged_file, merge_queue, merge_status],
        show_progress="hidden", api_name=False,
    )

    # Searching CivitAI, and dropping a result into a slot.
    search_btn.click(
        civitai_search, [search_query, search_base, search_nsfw],
        [search_results, search_pick, search_state], api_name=False,
    )
    search_query.submit(
        civitai_search, [search_query, search_base, search_nsfw],
        [search_results, search_pick, search_state], api_name=False,
    )
    search_put_btn.click(
        put_in_slot, [search_pick, search_state, search_slot],
        [*lora_references, search_results], 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,
    )

    # A CivitAI download link is a bare number, so the slots can be named from CivitAI's own public
    # model-versions endpoint: the title, the version, the file behind `fileId`, and the trigger words. Pressing
    # Enter in a slot names that set as well, so the button is only there for a paste that never gets an Enter.
    ir_button.click(
        build_ir_prompt,
        [prompt, ir_shot, ir_camera, ir_sound, ir_music, ir_speaker, ir_dialogue, ir_references],
        prompt,
        api_name=False,
    )

    lora_identify_btn.click(identify_loras, lora_references, lora_names, api_name=False)
    for _field in lora_references:
        _field.submit(identify_loras, lora_references, lora_names, 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)