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

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

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


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

# `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
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})"
    return (
        f"Ready · transformer + VAEs **bfloat16, unquantized** · placement `{PLACEMENT}` · attention "
        f"`{ATTENTION}` · {film} · {LORA_STATUS or 'no LoRA'} · loaded in {LOADED_IN:.0f}s · "
        f"conditioner `{CONDITIONER_SPACE}`"
    )


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

    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()
        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")
        pipe.load_components(dtype=torch.bfloat16)
        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")

        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

    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) * _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,
):
    """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
    with pk.use_schedule(
        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(PIPE, effective_seed, enabled=(sampler == "dpmpp_2s_ancestral")):
            with use_dpmpp_sde_gpu(PIPE, effective_seed, enabled=(sampler == "dpmpp_sde_gpu")):
                with use_seeds_2(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 {}
                    )
                    state = 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,
    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,
    progress=gr.Progress(track_tqdm=True),
    *,
    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)

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

    progress(0.1, desc=f"Denoising {this_stage_steps} steps at {width}x{height}, {num_frames} frames ...")
    call = (
        prompt_embeds,
        text_token_tags,
        keyframe(first_frame),
        keyframe(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,
    )
    # 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
    )
    return path, report, new_stage_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


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

            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)

    stage_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,
    )
    
    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,
        maximize_gpu,
        video_shift,
        audio_shift,
        sampler,
        stage_enabled,
        target_steps,
        stage_state,
        recondition,
    ]

    # `functools.partial` binds `advance` by keyword regardless of its position in `generate`'s signature — the
    # two buttons share every other line of conditioning/report logic and differ only in this one flag.
    run.click(
        functools.partial(generate, advance=False), controls, [video, report, stage_state], api_name="generate"
    )
    advance_btn.click(
        functools.partial(generate, advance=True), controls, [video, report, stage_state],
        api_name="generate_advance",
    )


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)