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 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")
LORA_FILES = {
    "lora1": os.environ.get("H3_LORA_1_FILE", "minimax_h3_turbo_v4_step600_ema.safetensors"),
    "loraa": os.environ.get("H3_LORA_A_FILE", "Mylo_lora_epoch31.safetensors"),
    "lorab": os.environ.get("H3_LORA_B_FILE", "VBVR_H3_attn_only.safetensors"),
    "lorac": os.environ.get("H3_LORA_C_FILE", "AIO_V2.safetensors"),
    "lorad": os.environ.get("H3_LORA_D_FILE", "Furry enhancer Video H3 V2.54.safetensors"),
    "lorae": os.environ.get("H3_LORA_E_FILE", "sb_H3_i2v_v1.1.safetensors"),
    "loraf": os.environ.get("H3_LORA_F_FILE", "moawxx_000002000.safetensors"),
    "lorag": os.environ.get("H3_LORA_G_FILE", "H3_ref2va_shot_v1_fp16.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": "Distilled LoRA",
    "loraa": "Anthro Enhancer",
    "lorab": "Reasoning Enhancer",
    "lorac": "HM-AIO", # hmmotion
    "lorad": "Anthro Realism",
    "lorae": "SB",
    "loraf": "Moaxx", # moawxx
    "lorag": "Fluid Enhancer",
}
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
# Some `diffusion_model.blocks.*` checkpoints store SwiGLU's fc1 gate/value halves in the opposite order
# diffusers expects. Leave off first; if the LoRA's effect looks inverted/broken rather than just weak or
# strong, set H3_LORA_SWAP_FC1=1 and compare.
SWAP_FC1_HALVES = os.environ.get("H3_LORA_SWAP_FC1", "0") == "1"

# 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


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 _convert_diffusion_model_lora(raw: dict, base_shapes: 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.
    """
    import re

    out = {}

    # Family A: `[diffusion_model.]blocks.N.(attn|mlp|adaln_proj).LEAF.(lora_A|lora_B).weight` — covers Mylo,
    # VBVR, AIO_V2, moawxx, the Furry Realism LoRA, and (minus its `diffusion_model.` prefix) the Turbo LoRA.
    standard = re.compile(
        r"^(?:diffusion_model\.)?blocks\.(\d+)\.(attn|mlp|adaln_proj)\.([\w.]+)\.(lora_[AB])\.weight$"
    )
    # Family B (Kohya-style): `lora_unet_blocks_N_TARGET.(lora_down|lora_up).weight` — covers SB and Fluid
    # Enhancer. `lora_down`/`lora_up` are the same A/B convention under a different name.
    kohya = re.compile(
        r"^lora_unet_blocks_(\d+)_(attn_out_proj|attn_qkv_proj|mlp_fc1|mlp_fc2)\.(lora_down|lora_up)\.weight$"
    )
    kohya_targets = {
        "attn_out_proj": ("attn", "out_proj"),
        "attn_qkv_proj": ("attn", "qkv_proj"),
        "mlp_fc1": ("mlp", "fc1"),
        "mlp_fc2": ("mlp", "fc2"),
    }
    kohya_ab = {"lora_down": "lora_A", "lora_up": "lora_B"}

    def emit(block: str, kind: str, leaf: str, ab: str, tensor) -> None:
        prefix = f"transformer_blocks.{block}."
        if kind == "attn" and leaf == "qkv_proj":
            if ab == "lora_A":
                # Shared low-rank input side — identical for q, k, v.
                out[f"{prefix}attn.to_q.{ab}.weight"] = tensor
                out[f"{prefix}attn.to_k.{ab}.weight"] = tensor
                out[f"{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"blocks.{block}.attn.qkv_proj.{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"{prefix}attn.to_q.{ab}.weight"] = tensor[:q_out].clone()
                out[f"{prefix}attn.to_k.{ab}.weight"] = tensor[q_out:q_out + k_out].clone()
                out[f"{prefix}attn.to_v.{ab}.weight"] = tensor[q_out + k_out:].clone()
        elif kind == "attn" and leaf == "out_proj":
            out[f"{prefix}attn.to_out.0.{ab}.weight"] = tensor
        elif kind == "mlp" and leaf == "fc1":
            if ab == "lora_B" and SWAP_FC1_HALVES:
                half = tensor.shape[0] // 2
                tensor = torch.cat([tensor[half:], tensor[:half]], dim=0)
            out[f"{prefix}ff.net.0.proj.{ab}.weight"] = tensor
        elif kind == "mlp" and leaf == "fc2":
            out[f"{prefix}ff.net.2.{ab}.weight"] = tensor
        elif kind == "adaln_proj" and leaf == "linear":
            out[f"{prefix}adaln_proj.linear.{ab}.weight"] = tensor
        else:
            print(f"[lora-convert] no mapping for blocks.{block}.{kind}.{leaf}.{ab}, skipping", flush=True)

    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 = standard.match(key)
        if match:
            block, kind, leaf, ab = match.groups()
            emit(block, kind, leaf, ab, tensor)
            continue

        match = kohya.match(key)
        if match:
            block, target, direction = match.groups()
            kind, leaf = kohya_targets[target]
            emit(block, kind, leaf, kohya_ab[direction], tensor)
            continue

        if key.endswith(".alpha"):
            # Per-module rank/alpha scaling isn't threaded through — matched modules get PEFT's default scaling
            # (scale 1.0), and the UI slider is what actually controls each LoRA's visible strength here. This
            # is a known simplification: a file's built-in alpha may have scaled it up or down from its raw
            # rank, so its slider range that "feels right" may not match what the file's author intended or
            # tested at. It isn't a bug — the Furry Realism LoRA's `.alpha` keys were already dropped the same
            # way and it loads and works fine — just worth knowing if a LoRA's effect seems unexpectedly
            # weak/strong across its whole slider range rather than at a specific value.
            continue

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

    return out

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."
    import h3_aoti

    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}` · {h3_aoti.status()} · {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 filename in LORA_FILES.values():
                try:
                    path = hf_hub_download(LORA_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, filename in LORA_FILES.items():
                try:
                    path = hf_hub_download(LORA_REPO, filename)
                    with safe_open(path, framework="pt") as handle:
                        raw = {k: handle.get_tensor(k) for k in handle.keys()}
                    converted = _convert_diffusion_model_lora(raw, base_shapes)
                    pipe.transformer.load_lora_adapter(converted, adapter_name=name, prefix=None)
                    # `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,
):
    """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:
        active = {name: strength for name, strength in lora_strengths.items() if name in LOADED_LORAS}
        if active:
            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()
    with pk.use_schedule(PIPE, steps, schedule, video_shift, audio_shift, sampler_name=sampler, seed=int(seed)):
        with use_dpmpp_2s_ancestral(PIPE, int(seed), enabled=(sampler == "dpmpp_2s_ancestral")):
            with use_dpmpp_sde_gpu(PIPE, int(seed), enabled=(sampler == "dpmpp_sde_gpu")):
                with use_seeds_2(PIPE, int(seed), enabled=(sampler == "seeds_2")):
                    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)),
                    )
            
    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")
    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


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_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,
    maximize_gpu=False,
    video_shift=DEFAULT_VIDEO_SHIFT,
    audio_shift=DEFAULT_AUDIO_SHIFT,
    sampler=DEFAULT_SAMPLER,
    progress=gr.Progress(track_tqdm=True),
):
    """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."""
    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)

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

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

    progress(0.1, desc=f"Denoising {int(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,
        int(steps),
        schedule_key,
        float(sharpen),
        multiplier,
        int(seed),
        lora_strengths,
        bool(maximize_gpu),
        float(video_shift),
        float(audio_shift),
        SAMPLERS.get(sampler, "euler"),
    )
    # 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 = _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()
    )
    report = (
        f"`{width}x{height}`, {num_frames} frames ({num_frames / FPS:.3f} s) -> {out_frames} frames at {fps} fps · "
        f"{int(steps)} steps of `{schedule_key}` · {' · '.join(post)} · seed {int(seed)} · {lora_text}\n\n"
        f"conditioner {condition_seconds:.0f}s ({plan['num_text_tokens']} tokens"
        f"{', upsampled' if refined else ''}) · denoise + decode {denoise_seconds:.0f}s "
        f"({denoise_seconds / max(1, int(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)
    return path, report


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)


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;}
"""

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")
            report = gr.Markdown()
            lora_1_strength = gr.Slider(
                label="Distilled LoRA",
                minimum=0.0,
                maximum=2.0,
                step=0.05,
                value=DEFAULT_LORA_1_STRENGTH,
            )
            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 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="Fluid Enhancer LoRA",
                    minimum=0.0,
                    maximum=2.0,
                    step=0.05,
                    value=DEFAULT_LORA_G_STRENGTH,
                )

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

    controls = [
        prompt,
        canvas,
        first_frame,
        last_frame,
        duration,
        steps,
        schedule,
        sharpen,
        interpolation,
        seed,
        upsample,
        lora_1_strength,
        lora_a_strength,
        lora_b_strength,
        lora_c_strength,
        lora_d_strength,
        lora_e_strength,
        lora_f_strength,
        lora_g_strength,
        maximize_gpu,
        video_shift,
        audio_shift,
        sampler,
    ]

    run.click(generate, controls, [video, report], api_name="generate")


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)