app.py

"""`Plaguekind/Minimax-H3` — the PlagueKind V1.5 ComfyUI workflow for MiniMax-H3, as a Space.
The candidate repository holds no weights: it is a ComfyUI graph over `Comfy-Org/MiniMax-H3`, so what is
reproduced here is the *graph*, on the `MiniMaxAI/MiniMax-H3` diffusers checkpoint. See `pk_workflow.py` for the
node-by-node mapping; the short version is euler + `linear_quadratic` at 15 steps, FSR RCAS sharpening at 0.3, and
FILM 2x frame interpolation to 48 fps.
Deployment is the split one the unquantized MiniMax-H3 needs: 195.9 GiB of bfloat16 does not fit under a Space's
150 GB storage quota, so the 62.14 GiB Qwen3-VL text encoder runs in a separate Space
(`multimodalart/qwen3vl-conditioner`) that this one calls per request, and this Space holds the 61.73 GiB
transformer and the two autoencoders. `prompt_embeds` + `text_token_tags` is the whole wire format.
"""

from __future__ import annotations

import functools
import os
import tempfile
import time
import traceback
from functools import cache

import torch

# 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 pk_workflow as pk
from h3_dpmpp_2s_ancestral import use_dpmpp_2s_ancestral, use_dpmpp_sde_gpu, use_seeds_2

MODEL_REPO = os.environ.get("H3_MODEL_REPO", "MiniMaxAI/MiniMax-H3")
CONDITIONER_SPACE = os.environ.get("H3_CONDITIONER", "dagloop5/qwen3vl-conditioner")
# `pack` places the transformer at startup, `lazy` moves everything on the first GPU call.
PLACEMENT = os.environ.get("H3_PLACEMENT", "pack").lower()
# cuDNN's fused attention is 10-20% faster than the SDPA default on this pool and needs nothing installed. It is
# also the closest available stand-in for the workflow's SageAttention patch, which is a sm90 build.
ATTENTION = os.environ.get("H3_ATTENTION", "_native_cudnn").lower()
GPU_SIZE = os.environ.get("H3_GPU_SIZE", "xlarge")

# A finetuned transformer, as a single monolithic safetensors file rather than MODEL_REPO's own sharded
# `transformer/` subfolder — everything else (VAE, schedulers, config) still comes from MODEL_REPO. Empty by
# default, which reproduces the official weights exactly. Confirmed diffusers-native key naming (not ComfyUI,
# not pruned-AdaLN) via `xal2077/PinkCherry_MiniMax-H3-Demo`'s own working `load_state_dict(strict=True)` call
# against an earlier release of the same lineage.
CUSTOM_TRANSFORMER_REPO = os.environ.get("H3_CUSTOM_TRANSFORMER_REPO", "SexGod1979/PinkCherry_MiniMax-H3")
CUSTOM_TRANSFORMER_FILE = os.environ.get(
    "H3_CUSTOM_TRANSFORMER_FILE", "v1-final-fl2va/PinkCherry_v1_bf16_fla2va_H3.safetensors"
)

LORA_REPO = os.environ.get("H3_LORA_REPO", "dagloop5/LoRA")
# Each entry is (repo, filename) so a LoRA can come from any repo, not just LORA_REPO — the two Lightx2v files
# live in lightx2v/Minimax-h3-Turbo, not dagloop5/LoRA.
LORA_FILES = {
    "lora1": (
        os.environ.get("H3_LORA_I_REPO", "alibaba-pai/MiniMax-H3-Acc-LoRAs"),
        os.environ.get("H3_LORA_I_FILE", "MiniMax-H3-FL2VA-Acc-8Step.safetensors"),
    ),
    "loraa": (LORA_REPO, os.environ.get("H3_LORA_A_FILE", "Mylo_lora_epoch31.safetensors")),
    "lorab": (LORA_REPO, os.environ.get("H3_LORA_B_FILE", "H3_VBVR_Pro_attn_only.safetensors")),
    "lorac": (LORA_REPO, os.environ.get("H3_LORA_C_FILE", "HM-AIO-V2.5.safetensors")),
    "lorad": (LORA_REPO, os.environ.get("H3_LORA_D_FILE", "Furry enhancer Video H3 V2.54.safetensors")),
    "lorae": (LORA_REPO, os.environ.get("H3_LORA_E_FILE", "sb_H3_i2v_v1.1.safetensors")),
    "loraf": (LORA_REPO, os.environ.get("H3_LORA_F_FILE", "moawxx_000002000.safetensors")),
    "lorag": (LORA_REPO, os.environ.get("H3_LORA_G_FILE", "Mystic_MMH3-V4.safetensors")),
    "lorah": (
        os.environ.get("H3_LORA_H_REPO", "lightx2v/Minimax-h3-Turbo"),
        os.environ.get("H3_LORA_H_FILE", "minimax_h3_fl2v_turbo_8step_v1.0_768p_comfyui_bf16.safetensors"),
    ),
    "lorai": (
        os.environ.get("H3_LORA_I_REPO", "lightx2v/Minimax-h3-Turbo"),
        os.environ.get("H3_LORA_I_FILE", "minimax_h3_fl2v_turbo_8step_v1.0_bf16.safetensors"),
    ),
    "loraj": (LORA_REPO, os.environ.get("H3_LORA_J_FILE", "H3_Motion_BoosterV2.safetensors")),
    "lorak": (LORA_REPO, os.environ.get("H3_LORA_k_FILE", "H3_Unlocked_V2.safetensors")),
    "loral": (LORA_REPO, os.environ.get("H3_LORA_l_FILE", "Ending_V1.safetensors")),
}
# Display names, keyed the same as LORA_FILES — used in the UI slider labels, the per-request report line, and
# the status line's failure list. Keep these two dicts' keys in sync when adding a LoRA.
LORA_LABELS = {
    "lora1": "MiniMax-H3-FL2VA-Acc-8Step",
    "loraa": "Anthro Enhancer",
    "lorab": "Reasoning Enhancer",
    "lorac": "HM-AIO", # hmmotion
    "lorad": "Anthro Realism",
    "lorae": "SB",
    "loraf": "Moaxx", # moawxx
    "lorag": "Mystic-V4",
    "lorah": "Lightx2v-Minimax-H3 Turbo 768p LoRA",
    "lorai": "Lightx2v-Minimax-H3 Turbo 8-step LoRA",
    "loraj": "Motion Booster V2",
    "lorak": "H3 Unlocked LoRA",
    "loral": "Ending LoRA",
}
DEFAULT_LORA_1_STRENGTH = 0.0
DEFAULT_LORA_A_STRENGTH = 0.0
DEFAULT_LORA_B_STRENGTH = 0.0
DEFAULT_LORA_C_STRENGTH = 0.0
DEFAULT_LORA_D_STRENGTH = 0.0
DEFAULT_LORA_E_STRENGTH = 0.0
DEFAULT_LORA_F_STRENGTH = 0.0
DEFAULT_LORA_G_STRENGTH = 0.0
DEFAULT_LORA_H_STRENGTH = 0.0
DEFAULT_LORA_I_STRENGTH = 0.0
DEFAULT_LORA_J_STRENGTH = 0.0
DEFAULT_LORA_K_STRENGTH = 0.0
DEFAULT_LORA_L_STRENGTH = 0.0
# Per-LoRA, not global: different training pipelines can store SwiGLU's fc1 gate/value halves in either order,
# and one flag can only be right for however many of the 8 files happen to agree. `lora1` (the Distilled/Turbo
# LoRA) is confirmed needing the swap by InstantX's official conversion of the same lineage
# (MiniMax-H3-Turbo-Lora-Diffusers/convert.py: "SwiGLU fc1 halves are swapped to match Diffusers' [value; gate]
# layout"); the rest default off until tested individually — set H3_LORA_SWAP_FC1_NAMES to a comma-separated
# list of LORA_FILES keys (e.g. "lora1,lorac") to override. Replaces H3_LORA_SWAP_FC1, which no longer does
# anything.
SWAP_FC1_NAMES = {name for name in os.environ.get("H3_LORA_SWAP_FC1_NAMES", "lora1").split(",") if name}

# 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. This is the workflow's "Target Dimension" node.
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),
}
# PlagueKind's V1.5 note: "FFLF is unreliable at res above 640". 960x544 keeps the short edge under that and is the
# canvas where the AoTI package pays most, so it is the default; the full 768 short edge is one dropdown away.
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.
MIN_UI_DURATION, MAX_UI_DURATION = 2, 30

SAMPLERS = {
    "euler": "euler",
    "euler ancestral": "euler_ancestral",
    "er_sde": "er_sde",
    "dpmpp_2m_sde_gpu": "dpmpp_2m_sde_gpu",
    "dpmpp_3m_sde_gpu": "dpmpp_3m_sde_gpu",
    "dpmpp_2s_ancestral": "dpmpp_2s_ancestral",
    "dpmpp_sde_gpu": "dpmpp_sde_gpu",
    "seeds_2": "seeds_2",
}
DEFAULT_SAMPLER = "euler"

SCHEDULES = {
    "linear_quadratic · PlagueKind": "linear_quadratic",
    "sgm_uniform": "sgm_uniform",
    "simple": "simple",
    "beta": "beta",
    "ddim_uniform": "ddim_uniform",
    "normal": "normal",
    "native (pipeline default)": "native",
}
DEFAULT_SCHEDULE = "linear_quadratic · PlagueKind"
# PlagueKind's original hardcoded values, now adjustable per request — the Turbo LoRA's own ComfyUI workflow
# uses video shift 6, not 12, so this is also how that gets tested against the Distilled LoRA.
DEFAULT_VIDEO_SHIFT = 12.0
DEFAULT_AUDIO_SHIFT = 3.0
INTERPOLATION = {"off · 24 fps": 1, "2x · 48 fps (PlagueKind)": 2, "4x · 96 fps": 4}
DEFAULT_INTERPOLATION = "2x · 48 fps (PlagueKind)"
DEFAULT_SHARPEN = 0.3
DEFAULT_STEPS = 15
# Staged Denoising: an arbitrary, adjustable starting point for the "Target total steps" slider — the total the
# fixed schedule is built at, walked across however many "Advance" presses it takes at "Steps" steps per press.
DEFAULT_TARGET_STEPS = 25
# Chunked Generation: an arbitrary, adjustable starting point for the "Chunk stop (s)" field.
DEFAULT_CHUNK_STOP = 10.0
# Momentum: seconds of the previous chunk's tail carried into the next chunk's opening. 0 disables momentum
# entirely, falling back to the plain last-frame-as-keyframe carry.
DEFAULT_MOMENTUM = 2.0


def snap_frames(seconds: float) -> int:
    """The frame count MiniMax-H3's video VAE can decode: the next `17 * n + 5` at 24 fps.
    Identical to the workflow's `ComfyMathExpression`,
    `max(5, round(a*24)) + (5 - (max(5, round(a*24)) % 17)) % 17` — 5 s is 124 frames, i.e. 5.167 s.
    """
    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_UI_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 raise_duration_ceiling(seconds: float = MAX_UI_DURATION) -> None:
    """Let the pipeline generate past its 15 s ceiling — experimental, past what MiniMax-H3 was trained/released
    at. `min_duration`/`max_duration` are the only place either bound is read (`before_denoise.py`'s validation
    step, `if not min_duration <= duration <= max_duration: raise ValueError`), and nothing architectural depends
    on the value: MiniMax-H3's RoPE computes `inv_freq` on the fly from arbitrary `position_ids`, not a
    fixed-size precomputed table, so there's no hard wraparound past 15 s — just untested territory, expect
    drift, looping, or identity loss rather than a clean extrapolation. Same technique as `lower_duration_floor`,
    the ceiling side.
    """
    from diffusers.modular_pipelines.minimax_h3.modular_pipeline import MiniMaxH3ModularPipeline

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

# `load_lora_adapter` requires every key (weights and `network_alphas` alike) to share a `prefix` whenever
# `network_alphas` is passed — `prefix=None` with a non-empty `network_alphas` is a hard error. This string is
# arbitrary (it's stripped off immediately, and the transformer itself has no `transformer.`-prefixed attribute)
# but has to match InstantX's own convention since it's just a filtering key, not a real path.
LORA_KEY_PREFIX = "transformer"

def _convert_diffusion_model_lora(raw: dict, base_shapes: dict, swap_fc1: bool) -> tuple[dict, dict]:
    """Rename a `diffusion_model.blocks.*` (original-checkpoint) LoRA state dict onto
    `MiniMaxH3Transformer3DModel`'s (`transformer_blocks.*`) naming, so `load_lora_adapter` can attach it.
    `raw` maps original key -> tensor. `base_shapes` maps the *unwrapped* base model's parameter names to their
    shapes — captured once before any adapter is attached, since `load_lora_adapter` wraps each target Linear in
    a PEFT layer and renames its weight to `<name>.base_layer.weight`, so a live `transformer.state_dict()` call
    after the first adapter attaches would no longer have `to_q.weight` etc. under their original names.

    Returns `(converted_weights, network_alphas)` — `network_alphas` is `load_lora_adapter`'s per-module `alpha`
    map. Built for every converted module, not just ones whose raw file carries an explicit `.alpha` key: PEFT's
    default scaling isn't guaranteed to land on `alpha == rank` when a LoRA mixes ranks across target types —
    InstantX's own Turbo-LoRA conversion needs `network_alphas` for exactly this reason (attn/mlp modules rank
    64, AdaLN modules rank 16), even though that file carries no `.alpha` keys at all. So `alpha = rank` is
    synthesized for every module first, then overridden wherever the raw file specifies something else.
    """
    import re

    out: dict = {}
    raw_alphas: dict[str, float] = {}  # raw ComfyUI base name -> alpha, from real `.alpha` keys only

    # Family A: standard (non-Kohya) naming — covers Mylo, VBVR, AIO_V2, moawxx, Anthro Realism, and (once
    # `diffusion_model.` is stripped) the Distilled/Turbo LoRA. Matched by module base name rather than one
    # fixed pattern, and renamed by substitution — this is what lets `token_refiner.blocks.*` and the top-level
    # `final_layer.adaln_proj` resolve onto real targets instead of falling through unmatched. Ported from
    # InstantX's official `MiniMax-H3-Turbo-Lora-Diffusers/convert.py`, written for this exact LoRA family.
    standard_ab = re.compile(r"^(?:diffusion_model\.)?(.+)\.(lora_[AB])\.weight$")
    standard_alpha = re.compile(r"^(?:diffusion_model\.)?(.+)\.alpha$")

    # Family C: already-diffusers-native, PEFT's own serialization layout — `{module}.lora_A.<adapter>.weight`,
    # confirmed against the debug dump's `transformer_blocks.0.*`/`token_refiner.refiner_blocks.*` shapes: no
    # fused `qkv_proj` to split (`to_q`/`to_k`/`to_v` are already separate), no `mlp.fc1`/`fc2` to rename (already
    # `ff.net.0.proj`/`ff.net.2`). The `<adapter>` segment is whatever adapter name the file happened to be saved
    # under (e.g. "default") — discarded, since each file gets its own `adapter_name` here regardless. No `.alpha`
    # keys exist in this family either (PEFT's native format keeps `lora_alpha` in a sidecar `adapter_config.json`
    # we never fetch, not as tensors), so these fall to the same `alpha = rank` default every other module gets.
    native_ab = re.compile(r"^(.+)\.(lora_[AB])\.\w+\.weight$")

    # Family B (Kohya-style): `lora_unet_blocks_N_TARGET.(lora_down|lora_up|alpha)` — covers SB and Fluid
    # Enhancer. `lora_down`/`lora_up` are the same A/B convention under a different name.
    kohya_ab = re.compile(
        r"^lora_unet_blocks_(\d+)_(attn_out_proj|attn_qkv_proj|mlp_fc1|mlp_fc2)\.(lora_down|lora_up)\.weight$"
    )
    kohya_alpha = re.compile(r"^lora_unet_blocks_(\d+)_(attn_out_proj|attn_qkv_proj|mlp_fc1|mlp_fc2)\.alpha$")
    kohya_targets = {
        "attn_out_proj": ("attn", "out_proj"),
        "attn_qkv_proj": ("attn", "qkv_proj"),
        "mlp_fc1": ("mlp", "fc1"),
        "mlp_fc2": ("mlp", "fc2"),
    }
    kohya_ab_name = {"lora_down": "lora_A", "lora_up": "lora_B"}

    def rename_base(name: str) -> str:
        """ComfyUI module path (before `.lora_*`/`.alpha`) -> Diffusers module path."""
        if name.startswith("token_refiner.blocks."):
            name = "token_refiner.refiner_blocks." + name[len("token_refiner.blocks."):]
        elif name.startswith("blocks."):
            name = "transformer_blocks." + name[len("blocks."):]
        name = name.replace("final_layer.adaln_proj.linear", "norm_out.linear")
        name = name.replace(".attn.out_proj", ".attn.to_out.0")
        name = name.replace(".mlp.fc2", ".ff.net.2")
        name = name.replace(".mlp.fc1", ".ff.net.0.proj")
        return name

    def target_bases(raw_base: str) -> list[str]:
        """Diffusers-side base name(s) for one pre-rename module path — one, except `attn.qkv_proj`, which fans
        out to `to_q`/`to_k`/`to_v` (same rank, so the same alpha applies to all three)."""
        if raw_base.endswith(".attn.qkv_proj"):
            prefix = rename_base(raw_base[: -len("attn.qkv_proj")])
            return [f"{prefix}attn.to_q", f"{prefix}attn.to_k", f"{prefix}attn.to_v"]
        return [rename_base(raw_base)]

    def emit(raw_base: str, ab: str, tensor) -> None:
        if raw_base.endswith(".attn.qkv_proj"):
            prefix = rename_base(raw_base[: -len("attn.qkv_proj")])
            if ab == "lora_A":
                # Shared low-rank input side — identical for q, k, v.
                out[f"{LORA_KEY_PREFIX}.{prefix}attn.to_q.{ab}.weight"] = tensor
                out[f"{LORA_KEY_PREFIX}.{prefix}attn.to_k.{ab}.weight"] = tensor
                out[f"{LORA_KEY_PREFIX}.{prefix}attn.to_v.{ab}.weight"] = tensor
            else:
                q_out = base_shapes[f"{prefix}attn.to_q.weight"][0]
                k_out = base_shapes[f"{prefix}attn.to_k.weight"][0]
                v_out = base_shapes[f"{prefix}attn.to_v.weight"][0]
                assert tensor.shape[0] == q_out + k_out + v_out, (
                    f"{raw_base}.{ab}: expected {q_out + k_out + v_out} rows "
                    f"(q{q_out}+k{k_out}+v{v_out}), got {tensor.shape[0]}"
                )
                out[f"{LORA_KEY_PREFIX}.{prefix}attn.to_q.{ab}.weight"] = tensor[:q_out].clone()
                out[f"{LORA_KEY_PREFIX}.{prefix}attn.to_k.{ab}.weight"] = tensor[q_out:q_out + k_out].clone()
                out[f"{LORA_KEY_PREFIX}.{prefix}attn.to_v.{ab}.weight"] = tensor[q_out + k_out:].clone()
            return

        if raw_base.endswith(".mlp.fc1") and ab == "lora_B" and swap_fc1:
            half = tensor.shape[0] // 2
            tensor = torch.cat([tensor[half:], tensor[:half]], dim=0)

        key_base = rename_base(raw_base)
        if f"{key_base}.weight" not in base_shapes:
            # The more permissive substitution-based rename can produce a name that isn't an actual target on
            # the live model — validated here rather than trusting the rename blindly, since it no longer
            # checks against a fixed whitelist of known `kind`s the way the old anchored regex did.
            print(f"[lora-convert] '{raw_base}' renamed to '{key_base}', which isn't a real target — skipping", flush=True)
            return
        out[f"{LORA_KEY_PREFIX}.{key_base}.{ab}.weight"] = tensor

    for key, raw_tensor in raw.items():
        # Some files (fp16-labeled ones especially) don't match the bf16 transformer's dtype; PEFT expects the
        # adapter's dtype to match the wrapped base layer's.
        tensor = raw_tensor.to(torch.bfloat16)

        match = native_ab.match(key)
        if match:
            module_base, ab = match.groups()
            if f"{module_base}.weight" in base_shapes:
                out[f"{LORA_KEY_PREFIX}.{module_base}.{ab}.weight"] = tensor
            else:
                print(f"[lora-convert] '{module_base}' isn't a real target — skipping", flush=True)
            continue

        match = standard_ab.match(key)
        if match:
            raw_base, ab = match.groups()
            emit(raw_base, ab, tensor)
            continue

        match = kohya_ab.match(key)
        if match:
            block, target, direction = match.groups()
            kind, leaf = kohya_targets[target]
            emit(f"blocks.{block}.{kind}.{leaf}", kohya_ab_name[direction], tensor)
            continue

        match = standard_alpha.match(key)
        if match:
            (raw_base,) = match.groups()
            raw_alphas[raw_base] = float(raw_tensor)
            continue

        match = kohya_alpha.match(key)
        if match:
            block, target = match.groups()
            kind, leaf = kohya_targets[target]
            raw_alphas[f"blocks.{block}.{kind}.{leaf}"] = float(raw_tensor)
            continue

        print(f"[lora-convert] skipping unrecognized key: {key}", flush=True)

    network_alphas: dict[str, float] = {}
    for out_key, out_tensor in out.items():
        if out_key.endswith(".lora_B.weight"):
            base = out_key[: -len(".lora_B.weight")]
            network_alphas[f"{base}.alpha"] = float(out_tensor.shape[1])
    for raw_base, alpha in raw_alphas.items():
        for base in target_bases(raw_base):
            network_alphas[f"{LORA_KEY_PREFIX}.{base}.alpha"] = alpha

    return out, network_alphas

PIPE = None
MOMENTUM_PIPE = None
MOMENTUM_ERROR: str | None = None
FILM = None
FILM_ERROR: str | None = None
LOAD_ERROR: str | None = None
LOADED_IN: float | None = None
LORA_STATUS: str | None = None
LOADED_LORAS: set[str] = set()


def status() -> str:
    if LOAD_ERROR:
        return LOAD_ERROR
    if PIPE is None:
        return f"Loading `{MODEL_REPO}` (transformer + VAEs, 77.3 GB). Watch the Space logs."

    film = "FILM **ready**" if FILM is not None else f"FILM **off** ({FILM_ERROR})"
    momentum = "momentum **ready**" if MOMENTUM_PIPE is not None else f"momentum **off** ({MOMENTUM_ERROR})"
    return (
        f"Ready · transformer + VAEs **bfloat16, unquantized** · placement `{PLACEMENT}` · attention "
        f"`{ATTENTION}` · {film} · {momentum} · {LORA_STATUS or 'no LoRA'} · loaded in {LOADED_IN:.0f}s · "
        f"conditioner `{CONDITIONER_SPACE}`"
    )

def _convert_full_checkpoint(raw: dict, base_shapes: dict) -> dict:
    """Rename a `diffusion_model.blocks.*`-family (original-checkpoint) full transformer state dict onto
    `MiniMaxH3Transformer3DModel`'s (`transformer_blocks.*`) naming — the full-weight sibling of
    `_convert_diffusion_model_lora`'s renaming: no A/B factors, no network_alphas, no fc1 swap, just every real
    weight tensor renamed (and, for the fused `qkv_proj`, split) onto its diffusers-side target. `base_shapes` is
    the target model's own shapes — safe to read straight off a `torch.device("meta")`-constructed instance,
    since shape is metadata, not data, and costs nothing to have before any real weights are loaded.
    """
    out: dict = {}

    def rename_base(name: str) -> str:
        if name.startswith("token_refiner.blocks."):
            name = "token_refiner.refiner_blocks." + name[len("token_refiner.blocks."):]
        elif name.startswith("blocks."):
            name = "transformer_blocks." + name[len("blocks."):]
        name = name.replace("final_layer.adaln_proj.linear", "norm_out.linear")
        name = name.replace("final_layer.norm", "norm_out.norm")
        name = name.replace("final_layer.video_out", "proj_out")
        name = name.replace("final_layer.audio_out", "audio_proj_out")
        name = name.replace(".attn.out_proj", ".attn.to_out.0")
        name = name.replace(".attn.q_norm", ".attn.norm_q")
        name = name.replace(".attn.k_norm", ".attn.norm_k")
        name = name.replace(".mlp.fc2", ".ff.net.2")
        name = name.replace(".mlp.fc1", ".ff.net.0.proj")
        name = name.replace("video_patch_proj", "proj_in")
        name = name.replace("audio_patch_proj", "audio_proj_in")
        name = name.replace("condition_proj", "context_embedder")
        name = name.replace("time_embedder.proj_in", "time_embedder.linear_1")
        name = name.replace("time_embedder.proj_out", "time_embedder.linear_2")
        return name

    for key, tensor in raw.items():
        if key == "rope.inv_freq":
            # A registered buffer computed from `rope_theta`/`rope_freq_dim` at construction, never loaded —
            # its presence here isn't a sign anything else is wrong.
            continue

        if key.endswith(".attn.qkv_proj.weight"):
            prefix = key[: -len(".attn.qkv_proj.weight")]
            renamed_prefix = rename_base(prefix)
            q_out = base_shapes[f"{renamed_prefix}.attn.to_q.weight"][0]
            k_out = base_shapes[f"{renamed_prefix}.attn.to_k.weight"][0]
            v_out = base_shapes[f"{renamed_prefix}.attn.to_v.weight"][0]
            assert tensor.shape[0] == q_out + k_out + v_out, (
                f"{key}: expected {q_out + k_out + v_out} rows (q{q_out}+k{k_out}+v{v_out}), got {tensor.shape[0]}"
            )
            out[f"{renamed_prefix}.attn.to_q.weight"] = tensor[:q_out].clone()
            out[f"{renamed_prefix}.attn.to_k.weight"] = tensor[q_out:q_out + k_out].clone()
            out[f"{renamed_prefix}.attn.to_v.weight"] = tensor[q_out + k_out:].clone()
            continue

        if key.endswith(".mlp.fc1.weight"):
            # Confirmed via `H3_DEBUG_COMPARE_ALL` against the official checkpoint: every one of the 52 fc1
            # layers (50 transformer blocks + 2 token-refiner blocks) mismatched, and nothing else did — the
            # exact signature of SwiGLU's gate/up halves being stored in the opposite order from what
            # `ff.net.0.proj` expects. The same swap `_convert_diffusion_model_lora` already applies for
            # `lora1`, here confirmed necessary for this repo's own full-checkpoint export.
            half = tensor.shape[0] // 2
            tensor = torch.cat([tensor[half:], tensor[:half]], dim=0)
            out[rename_base(key)] = tensor
            continue

        out[rename_base(key)] = tensor

    return out

def load_models() -> str | None:
    """Load the denoising half at startup, plus FILM.
    `MiniMaxH3GeneratorBlocks` declares `transformer`, `vae`, `audio_vae`, the two schedulers and `video_processor`,
    so `load_components` fetches exactly those subfolders — `text_encoder/` and `transformer_ref/` 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.
    """
    global PIPE, FILM, FILM_ERROR, LOAD_ERROR, LOADED_IN, LORA_STATUS, MOMENTUM_PIPE, MOMENTUM_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 MiniMaxH3GeneratorBlocks

        lower_duration_floor()
        raise_duration_ceiling()
        manager = ComponentsManager()
        blocks = MiniMaxH3GeneratorBlocks()
        print(f"[gen] 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")

        if CUSTOM_TRANSFORMER_REPO:
            # `load_config` fetches only `transformer/config.json` (a few KB) — not the 61.7 GiB of weights
            # `load_components` would otherwise pull from MODEL_REPO. Constructed on `torch.device("meta")` so
            # the architecture exists with no real memory behind it; `base_shapes` is read off that meta
            # instance (shape is metadata, not data) purely so `_convert_full_checkpoint` knows the real
            # to_q/to_k/to_v split points before any real weights exist. `load_state_dict(assign=True)`
            # materializes real tensors straight from the converted dict — the only real allocation here.
            from huggingface_hub import hf_hub_download
            from safetensors import safe_open

            from diffusers.models import MiniMaxH3Transformer3DModel

            config, _ = MiniMaxH3Transformer3DModel.load_config(
                MODEL_REPO, subfolder="transformer", return_unused_kwargs=True
            )
            with torch.device("meta"):
                custom_transformer = MiniMaxH3Transformer3DModel.from_config(config)
            base_shapes = {k: tuple(v.shape) for k, v in custom_transformer.state_dict().items()}

            custom_path = hf_hub_download(CUSTOM_TRANSFORMER_REPO, CUSTOM_TRANSFORMER_FILE)
            with safe_open(custom_path, framework="pt") as handle:
                raw = {k: handle.get_tensor(k) for k in handle.keys()}
            converted = _convert_full_checkpoint(raw, base_shapes)

            if os.environ.get("H3_DEBUG_COMPARE_ALL", "0") == "1":
                # The meta+assign mechanism is proven correct (H3_DEBUG_OFFICIAL_VIA_META produced coherent,
                # matching output using the official weights through this exact path), and two individual
                # tensors (`norm1.weight`, `attn.to_q.weight`) are already proven bit-exact — so the remaining
                # bug has to be in some *other* tensor `_convert_full_checkpoint` handles differently, not yet
                # individually checked. This downloads the official transformer once (same cost as the mechanism
                # test) and diffs every key against `converted`, rather than guessing which one to spot-check.
                import json as _json

                index_path = hf_hub_download(
                    MODEL_REPO, "transformer/diffusion_pytorch_model.safetensors.index.json"
                )
                with open(index_path) as handle:
                    index = _json.load(handle)
                shard_names = sorted(set(index["weight_map"].values()))
                official_state_dict = {}
                for shard_name in shard_names:
                    shard_path = hf_hub_download(MODEL_REPO, f"transformer/{shard_name}")
                    with safe_open(shard_path, framework="pt") as shard_handle:
                        for key in shard_handle.keys():
                            official_state_dict[key] = shard_handle.get_tensor(key)

                mismatches = []
                for key, official_tensor in official_state_dict.items():
                    converted_tensor = converted.get(key)
                    if converted_tensor is None:
                        mismatches.append((key, "missing from converted"))
                        continue
                    if converted_tensor.shape != official_tensor.shape:
                        mismatches.append((key, f"shape {tuple(converted_tensor.shape)} != {tuple(official_tensor.shape)}"))
                        continue
                    if not torch.allclose(converted_tensor.float(), official_tensor.float(), atol=1e-3):
                        diff = (converted_tensor.float() - official_tensor.float()).abs().max().item()
                        mismatches.append((key, f"values differ, max abs diff {diff:.6f}"))
                print(f"[debug-compare-all] checked {len(official_state_dict)} keys, {len(mismatches)} mismatches", flush=True)
                for key, reason in mismatches[:30]:
                    print(f"[debug-compare-all]   {key}: {reason}", flush=True)

            custom_transformer.load_state_dict(converted, strict=True, assign=True)

            # `rope.inv_freq` is a *non-persistent* buffer (`persistent=False` in `MiniMaxH3RotaryPosEmbed`) —
            # excluded from `state_dict()` entirely, which is exactly why `strict=True` above never complained
            # about its absence from the checkpoint (it was correctly dropped by `_convert_full_checkpoint` too).
            # But that also means `load_state_dict(assign=True)` never touches it: it's still sitting on
            # `torch.device("meta")` from construction, with nothing to move — which is what the later
            # `pipe.transformer.to("cuda")` call was hitting ("Cannot copy out of meta tensor; no data!").
            # Recomputed here from its own documented formula rather than moved, since a meta buffer has no data
            # to move in the first place.
            rope_theta = float(config["rope_theta"])
            rope_freq_dim = int(config["rope_freq_dim"])
            custom_transformer.rope.inv_freq = 1.0 / (
                rope_theta ** (torch.arange(0, 2 * rope_freq_dim, 2, dtype=torch.float32) / (2 * rope_freq_dim))
            )
            pipe.update_components(transformer=custom_transformer)
            print(f"[gen] transformer replaced with {CUSTOM_TRANSFORMER_REPO}/{CUSTOM_TRANSFORMER_FILE}", flush=True)

        pipe.load_components(dtype=torch.bfloat16)

        if CUSTOM_TRANSFORMER_REPO:
            # Moved to *after* `load_components(dtype=torch.bfloat16)` above, not before it: that call runs
            # unconditionally and its `dtype=` argument applies to every component regardless of whether it was
            # freshly fetched or already installed via `update_components` — doing this upcast beforehand had it
            # silently re-cast straight back to bf16 one line later, which is why the first attempt at this fix
            # had no visible effect at all despite being otherwise correct.
            #
            # `_keep_in_fp32_modules` is normally enforced by `from_pretrained`'s own post-load dtype pass — a
            # step this manual meta-device + `load_state_dict` path never goes through. The official checkpoint
            # genuinely ships these modules as float32 while everything else is bfloat16; if the finetune's file
            # is uniformly bf16 (common for a community single-file export), it adopted that dtype for these
            # layers too, silently dropping precision the model actually needs to run correctly.
            for name, tensor in list(pipe.transformer.named_parameters()) + list(pipe.transformer.named_buffers()):
                if any(keep in name for keep in MiniMaxH3Transformer3DModel._keep_in_fp32_modules):
                    tensor.data = tensor.data.to(torch.float32)

        pipe.transformer.set_attention_backend(ATTENTION)

        # --- Diagnostic: dump the LoRA files' key names/shapes and the transformer's own shapes to the Space
        # logs, so the exact rename map can be worked out without a notebook or shell. Set H3_LORA_DEBUG=0 in
        # the Space's env vars to silence this once you're done, or just delete this block later.
        if os.environ.get("H3_LORA_DEBUG", "1") != "0" and LORA_REPO.lower() not in ("", "off", "none"):
            from huggingface_hub import hf_hub_download
            from safetensors import safe_open

            for repo, filename in LORA_FILES.values():
                try:
                    path = hf_hub_download(repo, filename)
                    with safe_open(path, framework="pt") as handle:
                        keys = sorted(handle.keys())
                        print(f"[lora-debug] {filename}: {len(keys)} keys", flush=True)
                        for k in keys[:40]:
                            print(f"[lora-debug]   {k}  {tuple(handle.get_slice(k).get_shape())}", flush=True)
                        if len(keys) > 40:
                            print(f"[lora-debug]   ... and {len(keys) - 40} more", flush=True)
                except Exception as error:
                    print(f"[lora-debug] failed to inspect {filename}: {error}", flush=True)

            block0 = {
                k: tuple(v.shape)
                for k, v in pipe.transformer.state_dict().items()
                if k.startswith("transformer_blocks.0.")
            }
            print(f"[lora-debug] transformer_blocks.0.* ({len(block0)} keys):", flush=True)
            for k, shape in sorted(block0.items()):
                print(f"[lora-debug]   {k}  {shape}", flush=True)
            norm_out = {
                k: tuple(v.shape) for k, v in pipe.transformer.state_dict().items() if k.startswith("norm_out.")
            }
            print(f"[lora-debug] norm_out.* ({len(norm_out)} keys):", flush=True)
            for k, shape in sorted(norm_out.items()):
                print(f"[lora-debug]   {k}  {shape}", flush=True)

        # Approach B: convert each LoRA from its original `diffusion_model.blocks.*` naming onto this
        # transformer's `transformer_blocks.*` naming, then attach as PEFT layers, inactive (weight 0) until a
        # request asks for them. `load_lora_adapter` is the model-level loader (`PeftAdapterMixin`), used because
        # `MiniMaxH3ModularPipeline` has no pipeline-level `load_lora_weights` of its own.
        if LORA_REPO.lower() not in ("", "off", "none"):
            from huggingface_hub import hf_hub_download
            from peft.tuners.tuners_utils import BaseTunerLayer
            from safetensors import safe_open

            # Snapshot once, before any adapter attaches and wraps the target Linears — see the docstring on
            # `_convert_diffusion_model_lora` for why this can't be read fresh per-file.
            base_shapes = {k: tuple(v.shape) for k, v in pipe.transformer.state_dict().items()}
            failures = []
            for name, (repo, filename) in LORA_FILES.items():
                try:
                    path = hf_hub_download(repo, filename)
                    with safe_open(path, framework="pt") as handle:
                        raw = {k: handle.get_tensor(k) for k in handle.keys()}
                    converted, network_alphas = _convert_diffusion_model_lora(
                        raw, base_shapes, swap_fc1=name in SWAP_FC1_NAMES
                    )
                    pipe.transformer.load_lora_adapter(
                        converted, adapter_name=name, prefix=LORA_KEY_PREFIX, network_alphas=network_alphas
                    )
                    # `load_lora_adapter` warns-and-continues on a zero-key match instead of raising, so count
                    # matched layers ourselves and fail loudly if a file attached nothing.
                    matched = sum(
                        1
                        for module in pipe.transformer.modules()
                        if isinstance(module, BaseTunerLayer) and name in module.lora_A
                    )
                    if matched == 0:
                        raise RuntimeError(f"'{filename}' converted but matched 0 target modules")
                    LOADED_LORAS.add(name)
                except Exception as error:
                    failures.append(f"`{LORA_LABELS.get(name, name)}` ({type(error).__name__}: {error})")
                    print(
                        f"[gen] LoRA '{name}' ({filename}) failed to load: {type(error).__name__}: {error}",
                        flush=True,
                    )

            if LOADED_LORAS:
                pipe.transformer.set_adapters(list(LOADED_LORAS), weights=[0.0] * len(LOADED_LORAS))
            LORA_STATUS = "All LoRAs loaded" if not failures else "LoRA issues: " + "; ".join(failures)
            print(f"[gen] {LORA_STATUS}", flush=True)

        # Still startup, still free: an AoTI package carries no weights and opens its archive lazily inside the GPU
        # worker.
        import h3_aoti

        h3_aoti.maybe_load(pipe.transformer)

        if PLACEMENT == "pack":
            # Scoped to the transformer. `spaces` packs every startup-resident CUDA tensor into a second on-disk
            # copy, and packing all 77.3 GB busts the 150 GB storage quota; the 61.7 GB transformer alone fits. The
            # ~10 GB of fp32 VAEs move on the first GPU call instead.
            pipe.transformer.to("cuda")

        # Chunked Generation's momentum feature: a second, keyframe-free denoise graph over the *same* resident
        # weights — `update_components` links it to `pipe`'s own `transformer`/`vae`/`audio_vae`/schedulers, so
        # nothing here is loaded a second time. Soft-fail like FILM below: an experimental, newly-added block
        # (`h3_momentum.py`, untested end to end) failing here shouldn't take the whole Space down with it.
        try:
            from h3_momentum import MiniMaxH3MomentumGeneratorBlocks

            momentum_blocks = MiniMaxH3MomentumGeneratorBlocks()
            momentum_pipe = momentum_blocks.init_pipeline(MODEL_REPO, components_manager=manager, collection="h3")
            momentum_pipe.update_components(
                transformer=pipe.transformer,
                vae=pipe.vae,
                audio_vae=pipe.audio_vae,
                scheduler=pipe.scheduler,
                audio_scheduler=pipe.audio_scheduler,
            )
            MOMENTUM_PIPE = momentum_pipe
            print("[gen] momentum pipe ready", flush=True)
        except Exception as error:
            MOMENTUM_ERROR = f"{type(error).__name__}: {error}"
            print(f"[gen] momentum pipe unavailable ({MOMENTUM_ERROR}); momentum disabled", flush=True)

        PIPE = pipe
        LOADED_IN = time.time() - started
        print(f"[gen] ready in {LOADED_IN:.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

    # 69 MB of post-processing, and the demo is still a demo without it, so a failure here is not fatal.
    try:
        FILM = pk.load_film()
        print("[gen] FILM loaded", flush=True)
    except Exception as error:
        FILM_ERROR = f"{type(error).__name__}: {error}"
        print(f"[gen] FILM unavailable ({FILM_ERROR}); frame interpolation disabled", flush=True)

    return LOAD_ERROR


@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 encode_remote(prompt, image_path, last_image_path, canvas, num_frames, rewrite_prompt=False):
    """`/encode` 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."""
    from gradio_client import handle_file
    from safetensors import safe_open

    path, plan = conditioner().predict(
        prompt=prompt,
        image_path=handle_file(image_path) if image_path else None,
        last_image_path=handle_file(last_image_path) if last_image_path else None,
        canvas=canvas,
        num_frames=num_frames,
        rewrite_prompt=bool(rewrite_prompt),
        api_name="/encode",
    )
    with safe_open(path, framework="pt") as handle:
        metadata = handle.metadata()
        return handle.get_tensor("prompt_embeds"), handle.get_tensor("text_token_tags"), metadata, plan


# Seconds of GPU one request needs. Fitted to *this* Space against measurements, because booking a ceiling nobody
# reaches spends every visitor's ZeroGPU quota on nothing and costs the demo queue priority. Measured on the live
# Space: the default request takes 70 s and books 89; the first-and-last-frame one takes 79 s and books 94. The report
# each request prints carries both numbers, so the fit stays checkable.
#
# The denoise loop, from the packed video rows it is about to run: linear in the rows for the matmuls, quadratic for
# the attention, against the AoTI block package this Space loads. 3.6 s/step at the default canvas.
_DUR_B, _DUR_C = 1.1745e-4, 3.8396e-9
# The two resident decoders, which scale with the output rather than with the step count. `_DEFAULT_CANVAS_PIXELS` is
# 960x544x124, the default request, where the pair measures ~7 s.
_DECODE_BASE, _DECODE_PER_DEFAULT_CANVAS, _DEFAULT_CANVAS_PIXELS = 2, 5.5, 960 * 544 * 124
# The workflow's post chain. RCAS is a handful of elementwise passes over the clip; FILM is per *emitted* intermediate
# frame (a 2x pass over 124 frames is 123 of them); the h264 mux is per frame actually written.
_POST_BASE, _FILM_PER_FRAME, _MUX_PER_FRAME = 2.0, 0.025, 0.02
# `pack` mode: only the ~10 GB of fp32 VAEs move, and only on a cold worker.
_PLACEMENT_ALLOWANCE, _MARGIN = 8, 1.15
# The ZeroGPU per-call ceiling. A booking above it is refused with `ZeroGPU illegal duration` once the request is
# already in flight, so `generate` checks it up front and says which knob to turn instead.
_MAX_BOOKING = int(os.environ.get("H3_MAX_BOOKING", "1500"))
# Free-tier testing mode: forces the main Space's booking to exactly this many seconds regardless of the actual
# request. Paired with the conditioner Space's own fixed 8s booking (both xlarge), for a combined 148s against
# the shared 150s free-tier ceiling.
MAXIMIZE_GPU_DURATION = int(os.environ.get("H3_MAXIMIZE_GPU_DURATION", "140"))


def get_duration(
    prompt_embeds,
    text_token_tags,
    first_frame,
    last_frame,
    height,
    width,
    num_frames,
    steps,
    schedule,
    sharpen,
    multiplier,
    seed,
    lora_strengths,
    maximize_gpu,
    *a,
    **k,
):
    if maximize_gpu:
        return MAXIMIZE_GPU_DURATION

    # Momentum: `given_video` is the second-to-last item of `*a`, matching its fixed position at the end of the
    # `call` tuple in `generate()`. A flat, unvalidated allowance for the extra VAE encode — worth checking
    # against a real measurement once this is testable, the same as every other constant in this function was.
    given_video = a[-2] if len(a) >= 2 else None
    momentum_allowance = 5 if given_video is not None else 0

    height, width, num_frames, steps = int(height), int(width), int(num_frames), int(steps)
    multiplier = max(1, int(multiplier))
    latent_frames = (num_frames - LATENTS_PER_CHUNK) // FRAMES_PER_CHUNK * LATENTS_PER_CHUNK + 2
    patches = (height // 32) * (width // 32)
    keyframes = int(first_frame is not None) + int(last_frame is not None)
    rows = latent_frames * patches + keyframes * patches
    denoise = steps * (_DUR_B * rows + _DUR_C * rows**2)
    pixel_ratio = (height * width) / (960 * 544)
    decode = _DECODE_BASE + _DECODE_PER_DEFAULT_CANVAS * (height * width * num_frames) / _DEFAULT_CANVAS_PIXELS

    if multiplier > 1 and FILM is None:
        multiplier = 1
    out_frames = (num_frames - 1) * multiplier + 1 if multiplier > 1 else num_frames
    film = (num_frames - 1) * (multiplier - 1) * _FILM_PER_FRAME * pixel_ratio
    post = _POST_BASE + film + out_frames * _MUX_PER_FRAME * pixel_ratio

    return max(60, int((denoise + decode + post + momentum_allowance) * _MARGIN) + _PLACEMENT_ALLOWANCE)

@spaces.GPU(duration=get_duration, size=GPU_SIZE)
def _generate(
    prompt_embeds,
    text_token_tags,
    first_frame,
    last_frame,
    height,
    width,
    num_frames,
    steps,
    schedule,
    sharpen,
    multiplier,
    seed,
    lora_strengths,
    maximize_gpu,
    video_shift,
    audio_shift,
    sampler,
    total_steps,
    stage_from,
    resume_video_latents,
    resume_audio_latents,
    given_video,
    video_condition_mode,
):
    """The only thing on GPU time: the denoise loop, the two decoders and the workflow's post chain.
    The mp4 is muxed here rather than in the caller: a `@spaces.GPU` return crosses a process boundary by pickling,
    and a 2x-interpolated 124-frame clip is several hundred MB of frames against a few MB of h264.
    """
    import torch

    from diffusers.utils import encode_video

    global FILM

    booked = time.time()

# Approach B: blend whichever resident LoRA adapters actually loaded, for this request. Cheap —
    # `set_adapters` only updates each PEFT layer's active-adapter list and scale, no weight math — so it's safe
    # to call on every request. Filtered to `LOADED_LORAS`: a slider for a LoRA that failed at startup has no
    # adapter behind it, and `set_adapters` would raise if asked to activate a name that was never attached.
    if LOADED_LORAS:
        # Filtered to strength > 0, not just "loaded": PEFT computes every adapter in the active list on every
        # forward regardless of its weight (no early-exit for scale 0), so an adapter left active at 0.0 still
        # costs a real lora_A/lora_B matmul per targeted Linear, every block, every step — overhead that scales
        # with how many LoRAs are loaded, not how many are actually in use for a given request. Called
        # unconditionally, even with an empty list, rather than only `if active:` — skipping the call when every
        # slider is 0 would leave whichever adapters the *previous* request activated still live.
        active = {
            name: strength
            for name, strength in lora_strengths.items()
            if name in LOADED_LORAS and strength > 0
        }
        PIPE.transformer.set_adapters(list(active), weights=list(active.values()))

    if PLACEMENT == "lazy":
        PIPE.to("cuda")
    elif PLACEMENT == "pack":
        PIPE.vae.to("cuda")
        PIPE.audio_vae.to("cuda")

    steps = int(steps)
    multiplier = max(1, int(multiplier))
    custom_schedule = schedule != "native"
    # Any custom schedule — `linear_quadratic` or one of the five ported `BasicScheduler` names — hands
    # `set_timesteps` a finished `steps + 1` sigma grid, so it runs `steps` forwards. The native grid counts its
    # terminal zero as one of `num_inference_steps`, so it needs one more to match.
    requested_steps = steps if custom_schedule else steps + 1

    started = time.time()
    # A fresh generator per stage (see `use_schedule`'s docstring) is fine on its own — each stage's noise is
    # still a mathematically valid draw, just not a continuation of the last stage's stream. What isn't fine:
    # reseeding to the *identical* literal seed every single stage means every stage's first draws are the
    # exact same bits, every time — offsetting by `stage_from` (0 on a fresh/single-stage call, so this changes
    # nothing there) means each stage actually draws a different stream.
    effective_seed = int(seed) + stage_from
    # Momentum: a genuinely different, keyframe-free denoise graph, sharing every weight with `PIPE` (see
    # `load_models()`) — `pk.use_schedule`/the sampler context managers are already generic over whichever pipe
    # object they're handed, since both read and write the same shared `scheduler`/`audio_scheduler`/`transformer`.
    active_pipe = MOMENTUM_PIPE if given_video is not None else PIPE
    with pk.use_schedule(
        active_pipe, steps, schedule, video_shift, audio_shift, sampler_name=sampler, seed=effective_seed,
        total_steps=total_steps, stage_from=stage_from,
    ):
        with use_dpmpp_2s_ancestral(active_pipe, effective_seed, enabled=(sampler == "dpmpp_2s_ancestral")):
            with use_dpmpp_sde_gpu(active_pipe, effective_seed, enabled=(sampler == "dpmpp_sde_gpu")):
                with use_seeds_2(active_pipe, effective_seed, enabled=(sampler == "seeds_2")):
                    # Staged Denoising: resuming hands the pipeline the previous stage's own latents instead of
                    # letting `PrepareLatentsStep` draw fresh noise — both are declared-optional inputs on that
                    # step precisely for this ("used instead of the draw"), so nothing else about the call
                    # changes. `resume_video_latents is None` is exactly the unstaged, fresh-start case.
                    resume_kwargs = (
                        {"latents": resume_video_latents.to("cuda"), "audio_latents": resume_audio_latents.to("cuda")}
                        if resume_video_latents is not None
                        else {}
                    )
                    if given_video is not None:
                        # Momentum: keyframe-free, so no `image`/`last_image` at all — the carried clip already
                        # determines the opening frames more directly than a keyframe could.
                        state = active_pipe(
                            prompt_embeds=prompt_embeds.to("cuda"),
                            text_token_tags=text_token_tags,
                            height=height,
                            width=width,
                            num_frames=num_frames,
                            num_inference_steps=requested_steps,
                            output_type="pt",
                            generator=torch.Generator("cpu").manual_seed(int(seed)),
                            given_video=given_video.to("cuda"),
                            video_condition_mode=video_condition_mode,
                            **resume_kwargs,
                        )
                    else:
                        state = active_pipe(
                            prompt_embeds=prompt_embeds.to("cuda"),
                            text_token_tags=text_token_tags,
                            image=first_frame,
                            last_image=last_frame,
                            height=height,
                            width=width,
                            num_frames=num_frames,
                            num_inference_steps=requested_steps,
                            output_type="pt",
                            generator=torch.Generator("cpu").manual_seed(int(seed)),
                            **resume_kwargs,
                        )
            
    denoised = time.time() - started

    video = state.get("videos")[0]  # (frames, 3, H, W), float in [0, 1], on the card
    audio = state.get("audio")[0].cpu()
    sampling_rate = state.get("sampling_rate")
    # Staged Denoising: this stage's own final latents, ahead of decode — the state a later "Advance" press
    # resumes from. Computed unconditionally; harmless and cheap when staging isn't in use.
    stage_video_latents = state.get("latents").cpu()
    stage_audio_latents = state.get("audio_latents").cpu()
    del state
    # The post chain runs on the allocator the denoise loop just left fragmented (78.5 GiB at the full canvas), and
    # RCAS and FILM both want a few contiguous gigabytes.
    torch.cuda.empty_cache()

    post = time.time()
    video = pk.rcas(video, float(sharpen))
    if multiplier > 1:
        if FILM is None:
            multiplier = 1
        else:
            FILM = FILM.to("cuda")
            video = pk.interpolate(FILM, video, multiplier)
    fps = FPS * multiplier
    frames = (video.permute(0, 2, 3, 1).float() * 255.0).round_().clamp_(0, 255).to(torch.uint8).cpu()
    del video
    post_seconds = time.time() - post

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

    # `booked` to here is what `get_duration` had to predict, so it is what the report prints it against.
    return (
        path, denoised, post_seconds, time.time() - booked, int(frames.shape[0]), fps, multiplier,
        stage_video_latents, stage_audio_latents,
    )


def generate(
    prompt,
    canvas=DEFAULT_CANVAS,
    first_frame=None,
    last_frame=None,
    duration=5,
    steps=DEFAULT_STEPS,
    schedule=DEFAULT_SCHEDULE,
    sharpen=DEFAULT_SHARPEN,
    interpolation=DEFAULT_INTERPOLATION,
    seed=42,
    upsample=False,
    lora_1_strength=DEFAULT_LORA_1_STRENGTH,
    lora_h_strength=DEFAULT_LORA_H_STRENGTH,
    lora_i_strength=DEFAULT_LORA_I_STRENGTH,
    lora_a_strength=DEFAULT_LORA_A_STRENGTH,
    lora_b_strength=DEFAULT_LORA_B_STRENGTH,
    lora_c_strength=DEFAULT_LORA_C_STRENGTH,
    lora_d_strength=DEFAULT_LORA_D_STRENGTH,
    lora_e_strength=DEFAULT_LORA_E_STRENGTH,
    lora_f_strength=DEFAULT_LORA_F_STRENGTH,
    lora_g_strength=DEFAULT_LORA_G_STRENGTH,
    lora_j_strength=DEFAULT_LORA_J_STRENGTH,
    lora_k_strength=DEFAULT_LORA_K_STRENGTH,
    lora_l_strength=DEFAULT_LORA_L_STRENGTH,
    maximize_gpu=False,
    video_shift=DEFAULT_VIDEO_SHIFT,
    audio_shift=DEFAULT_AUDIO_SHIFT,
    sampler=DEFAULT_SAMPLER,
    stage_enabled=False,
    target_steps=DEFAULT_TARGET_STEPS,
    stage_state=None,
    recondition=True,
    chunk_enabled=False,
    chunk_start=0.0,
    chunk_stop=DEFAULT_CHUNK_STOP,
    chunk_state=None,
    momentum=DEFAULT_MOMENTUM,
    progress=gr.Progress(track_tqdm=True),
    *,
    advance: bool = False,
    chunk_advance: bool = False,
):
    """One request through the PlagueKind graph. Every parameter but the prompt carries the default its UI
    component carries, so an example that fills only `prompt` (and `canvas`) behaves exactly like the button.

    `advance` isn't a UI control — it's bound per-button via `functools.partial` (`False` for "Generate", `True`
    for "Advance") so the two share this one function rather than duplicating the conditioning/report logic.
    Staged Denoising, debugging-only, unlocked: nothing here stops the prompt, canvas, sampler, schedule, or
    shift from changing between an "Advance" press and the stage before it — the only samplers actually reasoned
    through for exact-vs-different-but-equal-quality resume behavior are `euler`, `euler_ancestral`, `seeds_2`,
    and `dpmpp_2s_ancestral`; the SDE-family samplers are untested here and not recommended.
    """
    if LOAD_ERROR:
        raise gr.Error(LOAD_ERROR)
    if PIPE is None:
        raise gr.Error("The denoiser is still loading.")
    if not prompt or not prompt.strip():
        raise gr.Error("MiniMax-H3 always takes a prompt, keyframes or not.")

    from PIL import Image, ImageOps

    canvas = canvas or DEFAULT_CANVAS
    schedule_key = SCHEDULES.get(schedule, "linear_quadratic")
    multiplier = INTERPOLATION.get(interpolation, 2)
    num_frames = snap_frames(duration)

    if stage_enabled and schedule_key == "native":
        raise gr.Error(
            "Staged Denoising needs a named sigma schedule, not `native` — the stage boundary is a slice of a "
            "schedule this Space builds itself, and the pipeline's own default schedule isn't one this Space "
            "controls the construction of."
        )
    if advance and stage_state is None:
        raise gr.Error("Press Generate with Staged Denoising enabled first, to start a staged sequence.")
    steps_done = int(stage_state["steps_done"]) if (advance and stage_state) else 0
    if advance:
        remaining = int(target_steps) - steps_done
        if remaining <= 0:
            raise gr.Error(
                f"Already at or past the target step count ({steps_done}/{int(target_steps)}). Raise "
                f"'Target total steps' to continue."
            )
        this_stage_steps = min(int(steps), remaining)
    else:
        this_stage_steps = int(steps)

    if chunk_enabled and stage_enabled:
        raise gr.Error("Chunked Generation and Staged Denoising can't both be enabled — pick one.")
    if chunk_advance and chunk_state is None:
        raise gr.Error("Press Generate with Chunked Generation enabled first, to start a chunked sequence.")

    chunk_first_frame = None
    given_video = None
    video_condition_mode = "locked"
    if chunk_advance:
        if MOMENTUM_PIPE is not None and float(momentum) > 0:
            given_video = _trailing_frames(chunk_state["paths"][-1], float(momentum))
        if given_video is None:
            # Momentum off, unavailable, or the extraction came back empty — fall back to the plain
            # last-frame-as-keyframe carry rather than dropping continuity entirely.
            chunk_first_frame = _last_frame_path(chunk_state["paths"][-1])
    effective_first_frame = chunk_first_frame if chunk_advance else first_frame
    effective_last_frame = None if chunk_enabled else last_frame

    if chunk_enabled:
        chunk_duration = float(chunk_stop) - float(chunk_start)
        if chunk_duration <= 0:
            raise gr.Error("Chunk stop must be after chunk start.")
        # A momentum-carrying chunk regenerates its own opening `momentum_seconds` from the previous chunk's
        # tail, which `_trim_head` removes again before appending — so the chunk's own generation has to run
        # `momentum_seconds` longer than requested, or trimming that regenerated span back off leaves less new
        # content than the chunk stop/start actually asked for. `snap_frames(momentum)/FPS`, not the raw slider
        # value, since that's the real, frame-aligned duration `_trailing_frames` actually extracted and
        # `MiniMaxH3MomentumConditionStep` actually imposed.
        momentum_seconds = snap_frames(float(momentum)) / FPS if given_video is not None else 0.0
        num_frames = snap_frames(chunk_duration + momentum_seconds)

    skip_recondition = advance and stage_state is not None and not recondition
    if skip_recondition:
        # "Re-condition" off: reuses this sequence's cached conditioning verbatim. Safe specifically because
        # nothing sampler/schedule/shift/steps/seed/sharpen/interpolation/LoRA-related is an input to the
        # conditioner at all — only prompt, the two keyframes, canvas, and "Upsample prompt" are. Height/width/
        # num_frames come from that same cached conditioning, so there's nothing new to compare for the
        # shape-consistency check below.
        prompt_embeds = stage_state["prompt_embeds"]
        text_token_tags = stage_state["text_token_tags"]
        metadata = stage_state["metadata"]
        plan = stage_state["plan"]
        condition_seconds = 0.0
        height, width, num_frames = stage_state["height"], stage_state["width"], stage_state["num_frames"]
        refined = stage_state.get("refined") or ""
    else:
        progress(
            0.0,
            desc=(
                f"Upsampling the prompt on {CONDITIONER_SPACE} ..."
                if upsample
                else f"Conditioning on {CONDITIONER_SPACE} ..."
            ),
        )
        conditioned = time.time()
        prompt_embeds, text_token_tags, metadata, plan = encode_remote(
            prompt, effective_first_frame, effective_last_frame, canvas, num_frames, rewrite_prompt=upsample
        )
        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 ""

        if advance and (height, width, num_frames) != (
            int(stage_state["height"]), int(stage_state["width"]), int(stage_state["num_frames"])
        ):
            raise gr.Error(
                "Canvas or duration resolved differently than the staged sequence's first stage — both have to "
                "stay fixed across a staged sequence, since they determine the saved latents' shape."
            )

    def keyframe(path):
        # The conditioning latents encoded here have to be of the image the conditioner looked at, which it
        # prepares exactly this way.
        return ImageOps.exif_transpose(Image.open(path)).convert("RGB") if path else None

    # Every UI LoRA slider gets packed into one dict here — this is the only place a new LoRA's slider value
    # needs wiring in; `_generate`, `set_adapters`, and the report line below are all keyed off this dict.
    lora_strengths = {"lora1": float(lora_1_strength), "lorah": float(lora_h_strength), "lorai": float(lora_i_strength), "loraa": float(lora_a_strength), "lorab": float(lora_b_strength), "lorac": float(lora_c_strength), "lorad": float(lora_d_strength), "lorae": float(lora_e_strength), "loraf": float(lora_f_strength), "lorag": float(lora_g_strength), "loraj": float(lora_j_strength), "lorak": float(lora_k_strength), "loral": float(lora_l_strength)}

    progress(0.1, desc=f"Denoising {this_stage_steps} steps at {width}x{height}, {num_frames} frames ...")
    call = (
        prompt_embeds,
        text_token_tags,
        keyframe(effective_first_frame),
        keyframe(effective_last_frame),
        height,
        width,
        num_frames,
        this_stage_steps,
        schedule_key,
        float(sharpen),
        multiplier,
        int(seed),
        lora_strengths,
        bool(maximize_gpu),
        float(video_shift),
        float(audio_shift),
        SAMPLERS.get(sampler, "euler"),
        int(target_steps) if stage_enabled else None,
        steps_done if advance else 0,
        stage_state["video_latents"] if advance else None,
        stage_state["audio_latents"] if advance else None,
        given_video,
        video_condition_mode,
    )
    # The same call `spaces` will book the worker with, so the report can show the fit against the measurement.
    booked_seconds = get_duration(*call)
    if booked_seconds > _MAX_BOOKING:
        raise gr.Error(
            f"That would book {booked_seconds}s of GPU, over the {_MAX_BOOKING}s ZeroGPU ceiling. Shorten the "
            f"**duration**, drop the **steps**, or pick a smaller **target dimension** — the denoise loop is "
            f"quadratic in the canvas."
        )
    (
        path, denoise_seconds, post_seconds, gpu_seconds, out_frames, fps, multiplier,
        stage_video_latents, stage_audio_latents,
    ) = _generate(*call)

    post = [f"RCAS {float(sharpen):.2f}" if float(sharpen) > 0 else "no sharpening"]
    post.append(f"FILM {multiplier}x -> {fps} fps" if multiplier > 1 else f"{fps} fps")
    lora_text = " / ".join(
        f"{LORA_LABELS.get(name, name)} {strength:.2f}"
        for name, strength in lora_strengths.items()
        if strength > 0
    )
    steps_done_after = steps_done + this_stage_steps
    info = [
        f"{this_stage_steps} steps of `{schedule_key}`",
        f"sampler `{sampler}`",
        f"shift {float(video_shift):.1f}/{float(audio_shift):.1f}",
        *post,
        f"seed {int(seed)}",
    ]
    if stage_enabled:
        info.append(f"staged {steps_done_after}/{int(target_steps)} steps")
    if lora_text:
        info.append(lora_text)
    report = (
        f"`{width}x{height}`, {num_frames} frames ({num_frames / FPS:.3f} s) -> {out_frames} frames at {fps} fps · "
        f"{' · '.join(info)}\n\n"
        f"conditioner {condition_seconds:.0f}s{' (cached)' if skip_recondition else ''} ({plan['num_text_tokens']} tokens"
        f"{', upsampled' if refined else ''}) · denoise + decode {denoise_seconds:.0f}s "
        f"({denoise_seconds / max(1, this_stage_steps):.1f} s/step) · post {post_seconds:.0f}s · "
        f"GPU {gpu_seconds:.0f}s of {booked_seconds}s booked"
    )
    if refined:
        report += f"\n\n**Upsampled prompt**\n\n{refined}"
    print(f"[gen] {report}", flush=True)

    new_stage_state = (
        {
            "video_latents": stage_video_latents,
            "audio_latents": stage_audio_latents,
            "height": height,
            "width": width,
            "num_frames": num_frames,
            "steps_done": steps_done_after,
            "prompt_embeds": prompt_embeds,
            "text_token_tags": text_token_tags,
            "metadata": metadata,
            "plan": plan,
            "refined": refined,
        }
        if stage_enabled
        else None
    )

    new_chunk_state = None
    if chunk_enabled:
        prior_paths = chunk_state["paths"] if chunk_advance else []
        # The momentum-imposed opening is a *regeneration* of the previous chunk's own tail, not new content —
        # trimmed here so concatenation doesn't duplicate it. Only continuation chunks that actually had momentum
        # applied carry anything to trim; chunk one, and any chunk that fell back to a plain keyframe carry, don't.
        chunk_output = _trim_head(path, momentum_seconds) if (chunk_advance and given_video is not None) else path
        chunk_paths = prior_paths + [chunk_output]
        path = _concat_chunks(chunk_paths) if len(chunk_paths) > 1 else chunk_output
        new_chunk_state = {"paths": chunk_paths}

    return path, report, new_stage_state, new_chunk_state


def _fit_keyframe(image_path, current_canvas):
    """Cover-crop an uploaded keyframe to the closest supported aspect ratio and select that ratio's smallest
    (fastest) canvas, unless the user already picked a matching ratio. The workflow's "Target Dimension" node does
    the same job by hand."""
    if not image_path:
        return gr.update(), gr.update()
    from PIL import Image as _Image

    img = _Image.open(image_path)
    aspect = img.width / img.height
    fastest = {}
    for label, (h, w) in CANVASES.items():
        r = w / h
        if r not in fastest or w * h < fastest[r][1][0] * fastest[r][1][1]:
            fastest[r] = (label, (h, w))
    ratio = min(fastest, key=lambda r: abs(r - aspect))
    label, (h, w) = fastest[ratio]

    cur_h, cur_w = CANVASES[current_canvas]
    if abs(cur_w / cur_h - aspect) <= abs(ratio - aspect):
        label = current_canvas
        h, w = cur_h, cur_w

    target = w / h
    if abs(img.width / img.height - target) <= 1e-3:
        return gr.update(), gr.update(value=label)
    if img.width / img.height > target:
        new_w = int(img.height * target)
        left = (img.width - new_w) // 2
        img = img.crop((left, 0, left + new_w, img.height))
    else:
        new_h = int(img.width / target)
        top = (img.height - new_h) // 2
        img = img.crop((0, top, img.width, top + new_h))
    img.save(image_path)
    return gr.update(value=image_path), gr.update(value=label)


# Client-side only (`fn=None`, no server round-trip): reads the real `<video>` element's current playback
# position, not a property of the file — so "grab this frame" means whatever's on screen when the button is
# pressed, paused or scrubbed to, not automatically the clip's last frame.
_FRAME_GRAB_JS = """
function() {
    const video = document.querySelector('#h3-generated-video video');
    return video ? video.currentTime : 0;
}
"""


def _extract_frame(video_path, timestamp):
    """The frame at `timestamp` seconds into `video_path`, as a numpy RGB array — Gradio converts it to a PIL
    image for whichever `gr.Image` this is wired to. Runs on CPU; no GPU time, no interaction with `_generate`."""
    if not video_path:
        return None
    import cv2

    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        return None
    fps = cap.get(cv2.CAP_PROP_FPS) or FPS
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    target_frame = min(int(float(timestamp) * fps), max(0, total_frames - 1))
    cap.set(cv2.CAP_PROP_POS_FRAMES, target_frame)
    ok, frame = cap.read()
    cap.release()
    return cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) if ok else None


def _last_frame_path(video_path: str) -> str | None:
    """The true last frame of `video_path`, saved as a temp PNG and handed back as a path — `_extract_frame`
    returns raw pixel data (built for populating a `gr.Image` component directly), and the keyframe path this
    feeds into `encode_remote`/`keyframe()` needs a real file, the same as an uploaded image would give one."""
    import cv2

    frame = _extract_frame(video_path, 1e9)  # 1e9 seconds: clamps to the true last frame
    if frame is None:
        return None
    directory = os.path.join(tempfile.gettempdir(), "pk-h3-chunk-heads")
    os.makedirs(directory, exist_ok=True)
    path = os.path.join(directory, f"chunk-head-{int(time.time() * 1000)}.png")
    cv2.imwrite(path, cv2.cvtColor(frame, cv2.COLOR_RGB2BGR))
    return path


def _trailing_frames(video_path: str, seconds: float):
    """The last `snap_frames(seconds)` of `video_path`'s pixel frames, at MiniMax-H3's own native 24 fps, as
    `(num_frames, 3, H, W)` **uint8** — `encode_vae_condition`'s own documented input convention (it does its
    own `/255` and ImageNet normalization internally; pre-dividing here would double it). The frame count is
    snapped to the same `17 * n + 5` the video VAE's temporal chunking requires for a multi-frame encode, per
    that function's own docstring — the same alignment `snap_frames` already gives a full request. Strided back
    to 24 fps first if the saved chunk was FILM-interpolated to a multiple of it: encoding frames at the wrong
    rate would encode the motion at the wrong speed. Runs on CPU; no GPU time.
    """
    if not video_path or seconds <= 0:
        return None
    import cv2
    import numpy as np
    import torch

    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        return None
    actual_fps = cap.get(cv2.CAP_PROP_FPS) or FPS
    stride = max(1, round(actual_fps / FPS))
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    target_frames = snap_frames(seconds)
    start_frame = max(0, total_frames - target_frames * stride)

    cap.set(cv2.CAP_PROP_POS_FRAMES, start_frame)
    frames = []
    for index in range(total_frames - start_frame):
        ok, frame = cap.read()
        if not ok:
            break
        if index % stride == 0:
            frames.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
        if len(frames) == target_frames:
            break
    cap.release()
    if len(frames) < target_frames:
        return None  # not enough source frames for a full, correctly-aligned encode
    array = np.stack(frames)
    return torch.from_numpy(array).permute(0, 3, 1, 2).contiguous()  # uint8, (num_frames, 3, H, W)


def _concat_chunks(paths: list[str]) -> str:
    """The chunks so far, concatenated with a stream copy (no re-encode) — cheap regardless of how many segments
    are in the list, so redoing the whole concat fresh on every press is simpler and more robust than trying to
    append onto an existing container."""
    import subprocess

    directory = os.path.join(tempfile.gettempdir(), "pk-h3-chunks")
    os.makedirs(directory, exist_ok=True)
    list_path = os.path.join(directory, f"concat-{int(time.time() * 1000)}.txt")
    with open(list_path, "w", encoding="utf-8") as handle:
        for path in paths:
            handle.write(f"file '{path}'\n")
    out_path = os.path.join(directory, f"chunked-{int(time.time() * 1000)}.mp4")
    subprocess.run(
        ["ffmpeg", "-y", "-f", "concat", "-safe", "0", "-i", list_path, "-c", "copy", out_path],
        check=True, capture_output=True,
    )
    return out_path

def _trim_head(video_path: str, seconds: float) -> str:
    """`video_path` with its first `seconds` cut off — the momentum-imposed opening a continuation chunk
    regenerates from the previous chunk's own tail, which would otherwise be duplicated once the chunks are
    concatenated. Re-encodes rather than stream-copying: an arbitrary, non-keyframe-aligned cut point can't
    always be trimmed losslessly with `-c copy`.
    """
    import subprocess

    directory = os.path.join(tempfile.gettempdir(), "pk-h3-chunks")
    os.makedirs(directory, exist_ok=True)
    out_path = os.path.join(directory, f"trimmed-{int(time.time() * 1000)}.mp4")
    subprocess.run(
        ["ffmpeg", "-y", "-ss", str(seconds), "-i", video_path, out_path],
        check=True, capture_output=True,
    )
    return out_path

load_models()

INTRO = """# PlagueKind · MiniMax-H3
<div align="center">
  <a href="https://huggingface.co/Plaguekind/Minimax-H3" target="_blank" rel="noopener"><strong>[ workflow ]</strong></a> &nbsp;
  <a href="https://huggingface.co/MiniMaxAI/MiniMax-H3" target="_blank" rel="noopener"><strong>[ model ]</strong></a> &nbsp;
  <a href="https://github.com/PlagueKind/Comfyui-PlagueKind-Nodes" target="_blank" rel="noopener"><strong>[ nodes ]</strong></a>
</div>
**MiniMax-H3** is a 33B parameter video generation model that produces video and a fully synchronized soundtrack
(ambience, foley, speech) in one pass. **PlagueKind's V1.5 workflow** is a tuning of it: euler on a
`linear_quadratic` sigma grid at 15 steps, FSR **RCAS** sharpening at 0.3, and **FILM** 2x frame interpolation to
48 fps. Text-to-video, first frame, last frame, or both.
"""

CSS = """
.main.fillable {max-width: 1250px !important}
.dark .gradio-container { color: var(--body-text-color); }
.status p {font-size: 0.8rem; opacity: 0.65; text-align: center;}
.h3-hidden-timestamp {
    opacity: 0;
    height: 0px;
    width: 0px;
    margin: 0px;
    padding: 0px;
    overflow: hidden;
    position: absolute;
    pointer-events: none;
}
"""

with gr.Blocks(title="PlagueKind · MiniMax-H3") as demo:
    gr.Markdown(INTRO)
    gr.Markdown(status(), elem_classes="status")

    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(
                label="Prompt",
                lines=3,
                value=(
                    "A red fox trotting through a snowy pine forest at dawn, snow crunching underfoot, "
                    "distant birdsong"
                ),
            )
            canvas = gr.Dropdown(
                label="Target dimension", choices=list(CANVASES), value=DEFAULT_CANVAS
            )
            with gr.Row():
                first_frame = gr.Image(label="First frame (optional)", type="filepath")
                last_frame = gr.Image(label="Last frame (optional)", type="filepath")
            run = gr.Button("Generate", variant="primary")

            with gr.Accordion("Advanced options", open=False):
                duration = gr.Slider(
                    label="Duration (s)",
                    minimum=MIN_UI_DURATION,
                    maximum=MAX_UI_DURATION,
                    step=1,
                    value=5,
                    info="Unstable past 15 seconds.",
                )
                steps = gr.Slider(
                    label="Steps",
                    minimum=4,
                    maximum=40,
                    step=1,
                    value=DEFAULT_STEPS,
                    info="PlagueKind: 15-20 on the linear_quadratic grid.",
                )
                sampler = gr.Dropdown(
                    label="Sampler",
                    choices=list(SAMPLERS),
                    value=DEFAULT_SAMPLER,
                    info="`euler ancestral` re-injects noise each step — expect seed to matter more.",
                )
                schedule = gr.Dropdown(
                    label="Sigma schedule",
                    choices=list(SCHEDULES),
                    value=DEFAULT_SCHEDULE,
                    info="`linear_quadratic` front-loads half the steps into the first 2.5% of the trajectory.",
                )
                video_shift = gr.Slider(
                    label="Video shift",
                    minimum=0.5,
                    maximum=50.0,
                    step=0.5,
                    value=DEFAULT_VIDEO_SHIFT,
                    info="Applies under every schedule, including native. LightX2V's Turbo LoRA uses 6, not 12.",
                )
                audio_shift = gr.Slider(
                    label="Audio shift",
                    minimum=0.5,
                    maximum=20.0,
                    step=0.5,
                    value=DEFAULT_AUDIO_SHIFT,
                )
                sharpen = gr.Slider(
                    label="RCAS sharpening",
                    minimum=0.0,
                    maximum=1.0,
                    step=0.05,
                    value=DEFAULT_SHARPEN,
                    info="FidelityFX Robust Contrast Adaptive Sharpening. PlagueKind: 0.3 is very natural.",
                )
                interpolation = gr.Dropdown(
                    label="FILM frame interpolation",
                    choices=list(INTERPOLATION),
                    value=DEFAULT_INTERPOLATION,
                    info="MiniMax-H3 generates 24 fps; FILM synthesizes the frames in between.",
                )
                seed = gr.Number(label="Seed", value=42, precision=0)
                upsample = gr.Checkbox(
                    label="Upsample prompt",
                    value=False,
                    info="Rewrite the prompt on the conditioner Space first, MiniMax's Context-IR style.",
                )
                maximize_gpu = gr.Checkbox(
                    label="Maximize Free Tier ZeroGPU (150 seconds)",
                    value=False,
                    info="Forces this request to book exactly 140s (plus 8s on the conditioner) for debugging purposes; does not prevent timeouts.",
                )

        with gr.Column():
            video = gr.Video(label="Video + soundtrack", elem_id="h3-generated-video")
            with gr.Row():
                grab_first_btn = gr.Button("📸 Use current frame as First frame", size="sm", variant="secondary")
                grab_last_btn = gr.Button("📸 Use current frame as Last frame", size="sm", variant="secondary")
            first_frame_timestamp = gr.Number(value=0, visible=True, elem_classes="h3-hidden-timestamp")
            last_frame_timestamp = gr.Number(value=0, visible=True, elem_classes="h3-hidden-timestamp")
            report = gr.Markdown()
            with gr.Accordion("Distilled / Turbo LoRAs", open=False):
                lora_1_strength = gr.Slider(
                    label="MiniMax-H3-FL2VA-Acc-8Step",
                    minimum=0.0,
                    maximum=2.0,
                    step=0.05,
                    value=DEFAULT_LORA_1_STRENGTH,
                    info="Video/Audio Shift = 6/3",
                )
                lora_h_strength = gr.Slider(
                    label="Lightx2v-Minimax-H3 Turbo 8-step 768p LoRA",
                    minimum=0.0,
                    maximum=2.0,
                    step=0.05,
                    value=DEFAULT_LORA_H_STRENGTH,
                    info="Video/Audio Shift = 6/3",
                )
                lora_i_strength = gr.Slider(
                    label="Lightx2v-Minimax-H3 Turbo 8-step LoRA",
                    minimum=0.0,
                    maximum=2.0,
                    step=0.05,
                    value=DEFAULT_LORA_I_STRENGTH,
                    visible=False,  # not confirmed working at its own shift yet — see the 8-step LoRA thread
                )
            with gr.Accordion("Custom LoRAs", open=False):
                lora_a_strength = gr.Slider(
                    label="Anthro Enhancer LoRA",
                    minimum=0.0,
                    maximum=2.0,
                    step=0.05,
                    value=DEFAULT_LORA_A_STRENGTH,
                )
                lora_b_strength = gr.Slider(
                    label="Reasoning Enhancer LoRA",
                    minimum=0.0,
                    maximum=2.0,
                    step=0.05,
                    value=DEFAULT_LORA_B_STRENGTH,
                )
                lora_c_strength = gr.Slider(
                    label="HM-AIO V2.5 LoRA",
                    minimum=0.0,
                    maximum=2.0,
                    step=0.05,
                    value=DEFAULT_LORA_C_STRENGTH,
                )
                lora_d_strength = gr.Slider(
                    label="Anthro Realism LoRA",
                    minimum=0.0,
                    maximum=2.0,
                    step=0.05,
                    value=DEFAULT_LORA_D_STRENGTH,
                )
                lora_e_strength = gr.Slider(
                    label="SB LoRA",
                    minimum=0.0,
                    maximum=2.0,
                    step=0.05,
                    value=DEFAULT_LORA_E_STRENGTH,
                )
                lora_f_strength = gr.Slider(
                    label="Moaxx LoRA",
                    minimum=0.0,
                    maximum=2.0,
                    step=0.05,
                    value=DEFAULT_LORA_F_STRENGTH,
                )
                lora_g_strength = gr.Slider(
                    label="Mystic V4.0 LoRA",
                    minimum=0.0,
                    maximum=2.0,
                    step=0.05,
                    value=DEFAULT_LORA_G_STRENGTH,
                )
                lora_j_strength = gr.Slider(
                    label="H3 Motion Booster V2 LoRA",
                    minimum=0.0,
                    maximum=2.0,
                    step=0.05,
                    value=DEFAULT_LORA_J_STRENGTH,
                )
                lora_k_strength = gr.Slider(
                    label="H3 Unlocked LoRA",
                    minimum=0.0,
                    maximum=2.0,
                    step=0.05,
                    value=DEFAULT_LORA_K_STRENGTH,
                )
                lora_l_strength = gr.Slider(
                    label="Ending LoRA",
                    minimum=0.0,
                    maximum=2.0,
                    step=0.05,
                    value=DEFAULT_LORA_L_STRENGTH,
                )

            with gr.Accordion("Staged Denoising", open=False):
                gr.Markdown(
                    "**Debugging feature — not for the SDE-family samplers** (`dpmpp_2m_sde_gpu`, "
                    "`dpmpp_3m_sde_gpu`, `dpmpp_sde_gpu`). Splits one long denoise into several cheaper requests: "
                    "run the first stage with **Generate**, then **Advance** to keep denoising the same latents "
                    "further, as many times as needed to reach the target."
                )
                stage_enabled = gr.Checkbox(label="Enable staged denoising", value=False)
                target_steps = gr.Slider(
                    label="Target total steps",
                    minimum=4,
                    maximum=100,
                    step=1,
                    value=DEFAULT_TARGET_STEPS,
                    visible=False,
                    info="The fixed schedule's total length — 'Steps' above is how many of these one press runs.",
                )
                recondition = gr.Checkbox(
                    label="Re-condition",
                    value=True,
                    visible=False,
                    info=(
                        "When turned off skips the conditioner on 'Advance' and reuses this sequence's cached prompt/keyframe "
                        "encoding — safe as long as the prompt, keyframes, target dimension, and 'Upsample prompt' haven't "
                        "changed since the first stage."
                    ),
                )
                advance_btn = gr.Button("Advance", variant="secondary", visible=False)

            with gr.Accordion("Chunked Generation", open=False):
                gr.Markdown(
                    "Splits a longer product into independent, full-quality chunks joined afterward — each "
                    "press is a complete generation, not a partial one. The previous chunk's last frame carries "
                    "into the next as its first frame."
                )
                chunk_enabled = gr.Checkbox(label="Enable chunked generation", value=False)
                chunk_start = gr.Number(label="Chunk start (s)", value=0.0, visible=False)
                chunk_stop = gr.Number(label="Chunk stop (s)", value=DEFAULT_CHUNK_STOP, visible=False)
                momentum = gr.Slider(
                    label="Momentum (s)",
                    minimum=0.0,
                    maximum=5.0,
                    step=0.5,
                    value=DEFAULT_MOMENTUM,
                    visible=False,
                    info=(
                        "Seconds of the previous chunk's tail imposed on the next chunk's opening, for real "
                        "motion continuity — 0 falls back to a plain last-frame keyframe. Untested past a "
                        "couple of seconds."
                    ),
                )
                continue_btn = gr.Button("Continue", variant="secondary", visible=False)

    stage_state = gr.State(None)
    chunk_state = gr.State(None)

    first_frame.upload(_fit_keyframe, [first_frame, canvas], [first_frame, canvas])
    last_frame.upload(_fit_keyframe, [last_frame, canvas], [last_frame, canvas])

    # Grabbing the currently-displayed frame: the button's own click runs only the JS above (`fn=None`, no
    # server round-trip) to read the real `<video>` element's playback position into a hidden number box; that
    # box's `.change()` is what actually decodes and writes the frame, server-side.
    grab_first_btn.click(fn=None, inputs=None, outputs=[first_frame_timestamp], js=_FRAME_GRAB_JS)
    first_frame_timestamp.change(
        _extract_frame, [video, first_frame_timestamp], first_frame, show_progress="hidden"
    )
    grab_last_btn.click(fn=None, inputs=None, outputs=[last_frame_timestamp], js=_FRAME_GRAB_JS)
    last_frame_timestamp.change(
        _extract_frame, [video, last_frame_timestamp], last_frame, show_progress="hidden"
    )

    stage_enabled.change(
        lambda enabled: tuple(gr.update(visible=enabled) for _ in range(3)),
        stage_enabled,
        [target_steps, recondition, advance_btn],
        api_name=False,
    )
    chunk_enabled.change(
        lambda enabled: tuple(gr.update(visible=enabled) for _ in range(4)),
        chunk_enabled,
        [chunk_start, chunk_stop, momentum, continue_btn],
        api_name=False,
    )

    controls = [
        prompt,
        canvas,
        first_frame,
        last_frame,
        duration,
        steps,
        schedule,
        sharpen,
        interpolation,
        seed,
        upsample,
        lora_1_strength,
        lora_h_strength,
        lora_i_strength,
        lora_a_strength,
        lora_b_strength,
        lora_c_strength,
        lora_d_strength,
        lora_e_strength,
        lora_f_strength,
        lora_g_strength,
        lora_j_strength,
        lora_k_strength,
        lora_l_strength,
        maximize_gpu,
        video_shift,
        audio_shift,
        sampler,
        stage_enabled,
        target_steps,
        stage_state,
        recondition,
        chunk_enabled,
        chunk_start,
        chunk_stop,
        chunk_state,
        momentum,
    ]

    # `functools.partial` binds `advance`/`chunk_advance` by keyword regardless of their position in `generate`'s
    # signature — the three buttons share every other line of conditioning/report logic and differ only in these
    # two flags.
    run.click(
        functools.partial(generate, advance=False, chunk_advance=False), controls,
        [video, report, stage_state, chunk_state], api_name="generate",
    )
    advance_btn.click(
        functools.partial(generate, advance=True, chunk_advance=False), controls,
        [video, report, stage_state, chunk_state], api_name="generate_advance",
    )
    continue_btn.click(
        functools.partial(generate, advance=False, chunk_advance=True), controls,
        [video, report, stage_state, chunk_state], api_name="generate_continue",
    ).then(
        lambda start, stop: (stop, min(stop + (stop - start), MAX_UI_DURATION * 100)),
        [chunk_start, chunk_stop], [chunk_start, chunk_stop],
        api_name=False,
    )


if __name__ == "__main__":
    # `theme` and `css` belong to `launch()` from Gradio 6.0 on; on `Blocks` they warn and are ignored.
    demo.launch(show_error=True, theme=gr.themes.Citrus(), css=CSS)