app.py

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

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

from __future__ import annotations

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

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

import pk_workflow as pk

MODEL_REPO = os.environ.get("H3_MODEL_REPO", "MiniMaxAI/MiniMax-H3")
CONDITIONER_SPACE = os.environ.get("H3_CONDITIONER", "dagloop5/qwen3vl-conditioner")
# `pack` moves only `transformer_ref` onto the card at startup, the same way fl2va scopes this to its own
# `transformer` partition: packing the full ~72.16 GiB pipe (plus `spaces`' on-disk pack copy) busts the 150 GB
# storage quota, but the ~61.7 GiB `transformer_ref` alone fits. The ~10 GB of fp32 VAEs move on the first GPU
# call instead. `lazy` moves everything on the first GPU call rather than packing anything; `offload` hands
# placement to `ComponentsManager.enable_auto_cpu_offload`.
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.
# flash-attention 3 is sm90-only and this card is sm120 (the `zero-a10g` flavour name is legacy).
ATTENTION = os.environ.get("H3_ATTENTION", "_native_cudnn").lower()
GPU_SIZE = os.environ.get("H3_GPU_SIZE", "xlarge")
# Bounds on what `get_duration` may reserve. The pool reserves whatever number it is given, so a flat ceiling for every
# request is what makes an account hit "too many ZeroGPU credits allocated to running tasks".
MIN_GPU_DURATION = int(os.environ.get("H3_GPU_DURATION_MIN", "120"))
MAX_GPU_DURATION = int(os.environ.get("H3_GPU_DURATION_MAX", "1500"))

# Must stay identical to the conditioner's table: the *label* goes over the wire, so a canvas that half does not know
# is rejected there and surfaces as a failure here.
CANVASES = {
    # 16:9
    "960x544 · 16:9 fast": (544, 960),
    "1024x576 · 16:9 fast": (576, 1024),
    "1152x640 · 16:9": (640, 1152),
    "1280x704 · 16:9": (704, 1280),
    "1344x768 · 16:9 full": (768, 1344),
    # 9:16
    "544x960 · 9:16 fast": (960, 544),
    "640x1152 · 9:16": (1152, 640),
    "768x1344 · 9:16 full": (1344, 768),
    # 1:1
    "544x544 · 1:1 fast": (544, 544),
    "768x768 · 1:1 full": (768, 768),
    # 4:3 / 3:4
    "768x576 · 4:3 fast": (576, 768),
    "1024x768 · 4:3 full": (768, 1024),
    "576x768 · 3:4 fast": (768, 576),
    "768x1024 · 3:4 full": (1024, 768),
    # 3:2 / 2:3
    "864x576 · 3:2 fast": (576, 864),
    "1152x768 · 3:2 full": (768, 1152),
    "576x864 · 2:3 fast": (864, 576),
    "768x1152 · 2:3 full": (1152, 768),
    # 21:9
    "1152x512 · 21:9 fast": (512, 1152),
    "1536x672 · 21:9 full": (672, 1536),
}
DEFAULT_CANVAS = "960x544 · 16:9 fast"

# Only the three plain scheduler.step() samplers from the fl2va Space — the SDE-family and two-evaluation
# samplers (dpmpp_2m/3m_sde_gpu, dpmpp_2s_ancestral, dpmpp_sde_gpu, seeds_2) aren't ported here.
SAMPLERS = {
    "euler": "euler",
    "euler ancestral": "euler_ancestral",
    "er_sde": "er_sde",
}
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"
DEFAULT_VIDEO_SHIFT = 12.0
DEFAULT_AUDIO_SHIFT = 3.0
DEFAULT_SHARPEN = 0.3
INTERPOLATION = {"off · 24 fps": 1, "2x · 48 fps (PlagueKind)": 2, "4x · 96 fps": 4}
DEFAULT_INTERPOLATION = "2x · 48 fps (PlagueKind)"

FPS, FRAMES_PER_CHUNK, LATENTS_PER_CHUNK = 24, 17, 5
# It is the *snapped* frame count the ceiling has to hold for: 15 s is 360 frames, which rounds up to 362, i.e.
# 15.083 s, and is refused. 14 is the last whole second that survives the snap.
MAX_UI_DURATION = 14
MIN_DURATION = 2
# A reference video shorter than 2 s gives the model almost no motion to read.
MIN_REFERENCE_VIDEO, MAX_REFERENCE_VIDEO = 2.0, 15.0
# `MINIMAX_H3_MAX_REFERENCE_IMAGES`. The slots are built up front and revealed one at a time, because a demo asking
# for two subjects should not open with nine boxes.
MAX_IMAGE_SLOTS, OPEN_IMAGE_SLOTS = 9, 2

MIN_STEPS = 4

# Seconds of GPU one request needs, from the packed sequence it is about to denoise: linear in the rows for the
# matmuls, quadratic for the attention, against the AoTI block package this Space runs.
STEP_LINEAR, STEP_QUADRATIC, SAFETY = 1.1745e-4, 3.8396e-9, 1.3
# `pack` mode: only the ~10 GB of fp32 VAEs move, and only on a cold worker — matches fl2va's own allowance for the
# identical move. Every request still carries it, because nothing here knows whether the worker it lands on is cold.
PLACEMENT_ALLOWANCE = int(os.environ.get("H3_PLACEMENT_ALLOWANCE", "8"))
AUDIO_LATENTS_PER_SECOND, AUDIO_CHANNELS = 40, 2
REFERENCE_IMAGE_SHORT_EDGE, CANVAS_MULTIPLE = 2048, 32
DECODE_BASE, DECODE_PER_DEFAULT_CANVAS, DEFAULT_CANVAS_PIXELS = 15, 25, 960 * 544 * 124
# The workflow's post chain. RCAS is a handful of elementwise passes over the clip; FILM is per *emitted*
# intermediate frame; the h264 mux is per frame actually written. Ported from the fl2va Space's fitted constants
# as a starting point, same caveat as `PLACEMENT_ALLOWANCE` above — worth checking booked-vs-actual here
# specifically once this is testable.
_POST_BASE, _FILM_PER_FRAME, _MUX_PER_FRAME = 2.0, 0.025, 0.02


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


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

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


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


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


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

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

    from diffusers.modular_pipelines.minimax_h3.modular_pipeline import resolve_canvas_size

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

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

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


def get_duration(
    prompt_embeds, text_token_tags, references, height, width, num_frames, steps, seed,
    sampler, schedule, video_shift, audio_shift, sharpen, multiplier, **_
):
    """Seconds of GPU to reserve for one request. Takes the arguments of the `@spaces.GPU` function it decorates, and
    tolerates the `gr.Progress` `spaces` injects."""
    sequence = int(text_token_tags.shape[0]) + reference_rows(references, num_frames) + target_rows(
        height, width, num_frames
    )
    denoise = int(steps) * (STEP_LINEAR * sequence + STEP_QUADRATIC * sequence**2) * SAFETY
    # The two reference encoders ahead of the loop, and the two decoders plus the mux after it. Both scale with what
    # they are handed rather than with the step count.
    encode = 5 + reference_rows(references, num_frames) * 1e-3
    decode = DECODE_BASE + DECODE_PER_DEFAULT_CANVAS * (height * width * num_frames) / DEFAULT_CANVAS_PIXELS

    multiplier = max(1, int(multiplier))
    if multiplier > 1 and FILM is None:
        multiplier = 1
    pixel_ratio = (height * width) / (960 * 544)
    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

    total = PLACEMENT_ALLOWANCE + encode + denoise + decode + post
    duration = max(MIN_GPU_DURATION, min(MAX_GPU_DURATION, int(total)))
    print(f"[ref2va] S={sequence} -> reserving {duration}s ({denoise:.0f}s of denoise at {steps} steps)", flush=True)
    return duration


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


def load_models() -> str | None:
    """Load the denoising half at startup, packing `transformer_ref` onto the card under `pack` placement.
    `MiniMaxH3Ref2VAGeneratorBlocks` declares `transformer_ref`, `vae`, `audio_vae`, the two schedulers and
    `video_processor`, so `load_components` fetches exactly those subfolders — `text_encoder/` and the `transformer/`
    partition are never touched. Both autoencoders carry `_keep_in_fp32_modules` over every module and stay float32: a
    bfloat16 audio VAE decodes the soundtrack roughly 20 dB too quiet.
    Only `transformer_ref` moves onto the card here, for storage rather than memory: `spaces`' startup `torch.pack()`
    writes every startup-resident CUDA tensor to a second copy on disk, and 77.3 GB of weights plus its pack busts
    the 150 GB quota (`OSError: [Errno 28] No space left on device` out of `os.posix_fallocate`, mid-pack); the
    ~61.7 GB `transformer_ref` alone fits, same as fl2va's `transformer`.
    """
    global PIPE, MANAGER, LOAD_ERROR, FILM, FILM_ERROR

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

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

        from h3_split_blocks import MiniMaxH3Ref2VAGeneratorBlocks

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

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

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

        h3_aoti.maybe_load(pipe.transformer_ref)

        if PLACEMENT == "offload":
            manager.enable_auto_cpu_offload(device="cuda")
            _arm_decode_hooks(pipe)
        elif PLACEMENT == "pack":
            # Scoped to `transformer_ref`, exactly as fl2va scopes this to its `transformer` partition — see the
            # docstring above for the quota math. The ~10 GB of fp32 VAEs move on the first GPU call instead.
            pipe.transformer_ref.to("cuda")

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

    # 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("[ref2va] FILM loaded", flush=True)
    except Exception as error:
        FILM_ERROR = f"{type(error).__name__}: {error}"
        print(f"[ref2va] FILM unavailable ({FILM_ERROR}); frame interpolation disabled", flush=True)

    return LOAD_ERROR


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

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

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

            setattr(module, method, armed)



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

    return Client(CONDITIONER_SPACE)


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

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

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


def _media_dimensions(path: str) -> tuple[int, int]:
    """`(width, height)` of an image or a video file, from its first stream."""
    from PIL import Image

    try:
        with Image.open(path) as image:
            return image.size
    except Exception:
        pass
    import av

    with av.open(path) as container:
        stream = container.streams.video[0]
        return stream.width, stream.height


def closest_canvas(path: str | None) -> str | None:
    """The canvas label whose aspect ratio is closest to a media file's, or `None` when the file is
    missing or unreadable."""
    if not path:
        return None
    try:
        width, height = _media_dimensions(path)
    except Exception:
        return None
    if not width or not height:
        return None
    target = width / height
    return min(CANVASES, key=lambda label: abs(CANVASES[label][1] / CANVASES[label][0] - target))


def auto_canvas(path):
    """Set the canvas to the closest aspect ratio of an uploaded image or video."""
    label = closest_canvas(path)
    return gr.update(value=label) if label else gr.update()


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

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


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

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


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


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


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


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

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

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


@spaces.GPU(duration=get_duration, size=GPU_SIZE)
def _generate(
    prompt_embeds, text_token_tags, references, height, width, num_frames, steps, seed,
    sampler, schedule, video_shift, audio_shift, sharpen, multiplier,
):
    """The only thing on GPU time: the two reference encoders, the packed-sequence denoise loop, the decoders, and
    the RCAS + FILM post chain. References cross as paths and are decoded here; only the generated outputs come
    back. A `@spaces.GPU` argument crosses a process boundary by pickling, a 5 s 1344x768 reference video is
    370 MB of expanded frames, and the full `PipelineState` still holds the packed latents and the rotary grid on
    the card.
    """
    import torch

    global FILM

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

    custom_schedule = schedule != "native"
    # Any custom schedule 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 = int(steps) if custom_schedule else int(steps) + 1

    with pk.use_schedule(PIPE, int(steps), schedule, video_shift, audio_shift, sampler_name=sampler, seed=int(seed)):
        state = PIPE(
            prompt_embeds=prompt_embeds.to("cuda"),
            text_token_tags=text_token_tags,
            references=build_references(references),
            height=height,
            width=width,
            num_frames=num_frames,
            num_inference_steps=requested_steps,
            output_type="pt",
            generator=torch.Generator("cpu").manual_seed(int(seed)),
        )

    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, and RCAS and FILM both want a
    # few contiguous gigabytes.
    torch.cuda.empty_cache()

    video = pk.rcas(video, float(sharpen))
    multiplier = max(1, int(multiplier))
    if multiplier > 1:
        if FILM is None:
            multiplier = 1
        else:
            FILM = FILM.to("cuda")
            video = pk.interpolate(FILM, video, multiplier)
    fps = FPS * multiplier

    # Muxed to an mp4 here, before returning, rather than in the caller: a raw CUDA tensor can't cross a
    # `@spaces.GPU` return at all under ZeroGPU's CUDA-emulation mode (`RuntimeError: Low-level CUDA init
    # reached` trying to reconstruct it in the dispatching process), and a CPU float tensor of several hundred
    # interpolated frames is needlessly large to pickle anyway when the finished file is a few MB of h264.
    from diffusers.utils import encode_video

    frames = (video.permute(0, 2, 3, 1).float() * 255.0).round_().clamp_(0, 255).to(torch.uint8).cpu()
    del video

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

    return path, fps, multiplier


def generate(
    prompt,
    image_1=None,
    audio_path=None,
    video_path=None,
    canvas=DEFAULT_CANVAS,
    image_2=None,
    image_3=None,
    image_4=None,
    image_5=None,
    image_6=None,
    image_7=None,
    image_8=None,
    image_9=None,
    match=True,
    duration=5,
    steps=28,
    seed=42,
    upsample=False,
    sampler=DEFAULT_SAMPLER,
    schedule=DEFAULT_SCHEDULE,
    video_shift=DEFAULT_VIDEO_SHIFT,
    audio_shift=DEFAULT_AUDIO_SHIFT,
    sharpen=DEFAULT_SHARPEN,
    interpolation=DEFAULT_INTERPOLATION,
    progress=gr.Progress(track_tqdm=True),
):
    """One request."""
    if LOAD_ERROR:
        raise gr.Error(LOAD_ERROR)
    if PIPE is None:
        raise gr.Error("The denoiser is still loading.")

    from diffusers.utils import encode_video

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

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

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

    progress(0.1, desc=f"Generating {num_frames / FPS:.1f} s at {width}x{height} ...")
    started = time.time()
    path, fps, multiplier = _generate(
        prompt_embeds, text_token_tags, references, height, width, num_frames, steps, seed,
        sampler_key, schedule_key, float(video_shift), float(audio_shift), float(sharpen), multiplier,
    )
    generate_seconds = time.time() - started

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

load_models()

INTRO = """# MiniMax-H3 Reference

<div align="center">
  <a href="https://huggingface.co/MiniMaxAI/MiniMax-H3" target="_blank" rel="noopener"><strong>[ model ]</strong></a> &nbsp;
  <a href="https://www.minimax.io/blog/minimax-h3" target="_blank" rel="noopener"><strong>[ blog ]</strong></a> &nbsp;
  <a href="https://huggingface.co/spaces/multimodalart/minimax-h3" target="_blank" rel="noopener"><strong>[ text / image to video ]</strong></a>
</div>

**MiniMax-H3** is a 33B parameter state of the art video generation model that produces video and a
fully synchronized soundtrack (ambience, foley, speech). Bring your own subject, voice or camera move as a
reference.
"""

CSS = """
.main.fillable { max-width: 1250px !important; }
.dark .gradio-container { color: var(--body-text-color); }
"""

with gr.Blocks(title="MiniMax-H3 Reference Custom Lora") as demo:
    gr.Markdown(INTRO)

    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(
                label="Prompt",
                lines=3,
                value="The character walks through a neon-lit street in the rain, humming to themselves",
            )
            upsample = gr.Checkbox(label="Upsample prompt", value=False)
            # One tab per modality, in the order the model reads them. A reference left in a tab that is not the open
            # one is still part of the request.
            with gr.Tabs():
                with gr.Tab("Images"):
                    # One `gr.Row`, so gradio splits the width evenly and wraps at `min_width` rather than leaving a
                    # hole where a hidden slot used to be.
                    with gr.Row():
                        images = [
                            gr.Image(
                                label="Subject, style or scene",
                                type="filepath",
                                min_width=180,
                                # Fixed, so a row that wraps to a single slot stays the size of a full one.
                                height=210,
                                visible=index < OPEN_IMAGE_SLOTS,
                            )
                            for index in range(MAX_IMAGE_SLOTS)
                        ]
                    add_image = gr.Button("+ Add another image", size="sm", variant="secondary")
                with gr.Tab("Audio"):
                    audio = gr.Audio(label="A voice or a piece of music", type="filepath")
                with gr.Tab("Video"):
                    video = gr.Video(label="Motion & camera, 2–15 s. Its soundtrack comes along.")
            run = gr.Button("Generate", variant="primary")
            with gr.Accordion("Advanced options", open=False):
                canvas = gr.Dropdown(label="Canvas", choices=list(CANVASES), value=DEFAULT_CANVAS)
                match = gr.Checkbox(label="Match the reference soundtrack", value=True, visible=False)
                duration = gr.Slider(
                    label="Duration (s)", minimum=MIN_DURATION, maximum=MAX_UI_DURATION, step=1, value=5
                )
                steps = gr.Slider(label="Steps", minimum=MIN_STEPS, maximum=40, step=1, value=28)
                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
                )
                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)

        with gr.Column():
            result = gr.Video(label="Video + soundtrack")
            # An output, so it can be revealed only for a request that asked for a rewrite.
            with gr.Accordion("Upsampled prompt", open=False, visible=False) as upsampled_panel:
                upsampled = gr.Textbox(show_label=False, lines=8, interactive=False)

    open_slots = gr.State(OPEN_IMAGE_SLOTS)

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

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

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

    # Auto-select the canvas whose aspect ratio is closest to an uploaded image or video.
    for image in images:
        image.change(auto_canvas, image, canvas, show_progress="hidden", api_name=False)
    video.change(auto_canvas, video, canvas, show_progress="hidden", api_name=False)

    request = [
        prompt, images[0], audio, video, canvas, *images[1:], match, duration, steps, seed, upsample,
        sampler, schedule, video_shift, audio_shift, sharpen, interpolation,
    ]
    run.click(generate, request, [result, upsampled, upsampled_panel], api_name="generate")


if __name__ == "__main__":
    demo.launch(show_error=True, theme=gr.themes.Citrus(), css=CSS)