h3_engine.py

"""The denoising half of MiniMax-H3, plus everything a request needs around it.

MiniMax-H3 is 196 GiB in bfloat16 and a ZeroGPU Space is evicted past 150 GB of storage, so the pipeline is split at
its text-encoder step: the 62 GiB Qwen3-VL conditioner runs in a Space of its own (`H3_CONDITIONER`), this module
holds the transformer and the two autoencoders, and `prompt_embeds` + `text_token_tags` is the whole wire format
between the two halves.

Nothing here inspects, rewrites or refuses a prompt: the string is handed to the conditioner as typed.
"""

from __future__ import annotations

import hashlib
import os
import tempfile
import time
import traceback
from collections import OrderedDict
from functools import cache

# `import spaces` patches `torch.cuda` before anything can initialize it, which is what lets the 77 GB load happen at
# startup instead of on booked GPU time. It has to stay the first import that touches torch. Off the ZeroGPU runtime
# (a laptop running `DEV_MODE=1`) the package is absent and the decorator is a no-op.
try:
    import spaces
except ModuleNotFoundError:
    from types import SimpleNamespace

    spaces = SimpleNamespace(GPU=lambda *args, **kwargs: (lambda function: function))

import lora_stack

MODEL_REPO = os.getenv("H3_MODEL_REPO", "MiniMaxAI/MiniMax-H3")
CONDITIONER_SPACE = os.getenv("H3_CONDITIONER", "multimodalart/qwen3vl-conditioner")
GPU_SIZE = os.getenv("H3_GPU_SIZE", "xlarge")
# cuDNN's fused attention needs nothing installed and beats the SDPA default on this pool.
ATTENTION = os.getenv("H3_ATTENTION", "_native_cudnn")
# "startup" keeps the transformer resident from boot; "lazy" moves the whole pipeline on the first request instead.
PLACEMENT = os.getenv("H3_PLACEMENT", "startup").lower()
# The released checkpoint refuses anything under 5 s. It generates fine well below that, so the floor is a knob.
MIN_DURATION = float(os.getenv("H3_MIN_DURATION", "2"))
DEV_MODE = os.getenv("DEV_MODE", "0") == "1"

FPS = 24
# The video VAE decodes `17 * n + 5` frames and nothing else.
FRAMES_PER_CHUNK, LATENTS_PER_CHUNK = 17, 5
MIN_UI_DURATION, MAX_UI_DURATION = 2, 14

# The canvas labels are the conditioner's own table - the *label* goes over the wire, so a canvas this side invents
# is rejected there. Verified against `/encode`'s enum.
CANVASES: dict[str, tuple[int, int]] = {
    "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),
    "544x960 · 9:16 fast": (960, 544),
    "640x1152 · 9:16": (1152, 640),
    "768x1344 · 9:16 full": (1344, 768),
    "544x544 · 1:1 fast": (544, 544),
    "768x768 · 1:1 full": (768, 768),
    "1024x1024 · 1:1 max": (1024, 1024),
    "768x576 · 4:3 fast": (576, 768),
    "1024x768 · 4:3 full": (768, 1024),
    "576x768 · 3:4 fast": (768, 576),
    "768x1024 · 3:4 full": (1024, 768),
    "1152x512 · 21:9 fast": (512, 1152),
    "1536x672 · 21:9 full": (672, 1536),
}
DEFAULT_CANVAS = "960x544 · 16:9 fast"

OUTPUT_DIR = os.path.join(tempfile.gettempdir(), "h3-clips")

# Conditionings kept in memory, newest last. A few megabytes each; this is the main process, shared by everyone
# using the Space, and the conditioner is a pure function of its inputs, so a shared entry is the same answer.
ENCODE_CACHE_SIZE = int(os.getenv("H3_ENCODE_CACHE", "32"))
_encode_cache: OrderedDict[str, tuple] = OrderedDict()

PIPE = None
LOAD_ERROR: str | None = None
LOADED_IN: float | None = None


# ---------------------------------------------------------------------------
# Geometry
# ---------------------------------------------------------------------------
def snap_frames(seconds: float) -> int:
    """The frame count the 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 latent_frames(num_frames: int) -> int:
    """Latent frames the video VAE produces for an aligned frame count: `5 * n + 2`."""
    return (num_frames - LATENTS_PER_CHUNK) // FRAMES_PER_CHUNK * LATENTS_PER_CHUNK + 2


def fastest_canvas_per_ratio() -> dict[float, str]:
    """The cheapest canvas of every supported aspect ratio, keyed by that ratio."""
    best: dict[float, tuple[str, int]] = {}
    for label, (height, width) in CANVASES.items():
        ratio = width / height
        if ratio not in best or width * height < best[ratio][1]:
            best[ratio] = (label, width * height)
    return {ratio: label for ratio, (label, _) in best.items()}


def fit_keyframe(path: str, canvas: str) -> tuple[str, str]:
    """Cover-crop an uploaded keyframe onto the closest supported aspect ratio, in place.

    The picked canvas is the cheapest one of that ratio, unless the canvas the user chose is already at least as
    close to the image - a deliberate 21:9 pick survives a 16:9-ish upload. Returns `(path, canvas_label)`.
    """
    from PIL import Image, ImageOps

    image = ImageOps.exif_transpose(Image.open(path)).convert("RGB")
    aspect = image.width / image.height

    candidates = fastest_canvas_per_ratio()
    ratio = min(candidates, key=lambda r: abs(r - aspect))
    label = candidates[ratio]

    chosen_h, chosen_w = CANVASES[canvas]
    if abs(chosen_w / chosen_h - aspect) <= abs(ratio - aspect):
        label = canvas
    height, width = CANVASES[label]

    target = width / height
    if abs(aspect - target) > 1e-3:
        if aspect > target:
            crop_w = round(image.height * target)
            left = (image.width - crop_w) // 2
            image = image.crop((left, 0, left + crop_w, image.height))
        else:
            crop_h = round(image.width / target)
            top = (image.height - crop_h) // 2
            image = image.crop((0, top, image.width, top + crop_h))

    image.save(path)
    return path, label


# ---------------------------------------------------------------------------
# Loading
# ---------------------------------------------------------------------------
def _generator_blocks():
    """The `t2va` / `fl2va` half of `MiniMaxH3Blocks` with the text-encoder step cut out.

    The `ref2va` branches are left out on purpose: their denoise step declares `transformer_ref`, and declaring it
    would make `load_components` pull a second 61.7 GB partition this Space has no room for.
    """
    import torch
    from diffusers.modular_pipelines.minimax_h3.before_encoder import MiniMaxH3ResizeStep
    from diffusers.modular_pipelines.minimax_h3.encoders import MiniMaxH3KeyframeVaeEncoderStep
    from diffusers.modular_pipelines.minimax_h3.modular_blocks_minimax_h3 import (
        MiniMaxH3CoreDenoiseStep,
        MiniMaxH3DecodeStep,
        MiniMaxH3FL2VACoreDenoiseStep,
    )
    from diffusers.modular_pipelines.modular_pipeline import ConditionalPipelineBlocks, SequentialPipelineBlocks
    from diffusers.modular_pipelines.modular_pipeline_utils import OutputParam

    def keyframed(kwargs) -> bool:
        return kwargs.get("image") is not None or kwargs.get("last_image") is not None

    class KeyframeEncodeStep(ConditionalPipelineBlocks):
        model_name = "minimax-h3"
        block_classes = [MiniMaxH3KeyframeVaeEncoderStep]
        block_names = ["keyframes"]
        block_trigger_inputs = ["image", "last_image"]
        default_block_name = None

        def select_block(self, **kwargs) -> str | None:
            return "keyframes" if keyframed(kwargs) else None

        @property
        def description(self):
            return "Encodes the keyframes into conditioning latents; a text-only request skips it."

    class DenoiseStep(ConditionalPipelineBlocks):
        model_name = "minimax-h3"
        block_classes = [MiniMaxH3FL2VACoreDenoiseStep, MiniMaxH3CoreDenoiseStep]
        block_names = ["fl2va", "t2va"]
        block_trigger_inputs = ["image", "last_image"]
        default_block_name = "t2va"

        def select_block(self, **kwargs) -> str | None:
            return "fl2va" if keyframed(kwargs) else None

        @property
        def description(self):
            return "The packed-sequence denoise loop, on the `fl2va` branch when a keyframe was passed."

    class GeneratorBlocks(SequentialPipelineBlocks):
        """Everything after the conditioner: keyframes onto the canvas, denoise, decode."""

        model_name = "minimax-h3"
        block_classes = [MiniMaxH3ResizeStep, KeyframeEncodeStep, DenoiseStep, MiniMaxH3DecodeStep]
        block_names = ["resize", "vae_encoder", "denoise", "decode"]

        @property
        def description(self):
            return (
                "The denoising half of a split MiniMax-H3 deployment: `prompt_embeds` and `text_token_tags` arrive "
                "as inputs, so the 62 GiB Qwen3-VL text encoder is never loaded here."
            )

        @property
        def outputs(self):
            return [
                OutputParam.template("videos", description="The generated video."),
                OutputParam("audio", type_hint=torch.Tensor, description="The generated stereo soundtrack."),
                OutputParam("sampling_rate", type_hint=int, description="Sample rate of the soundtrack, in Hz."),
            ]

    return GeneratorBlocks()


def _lower_duration_floor(seconds: float) -> None:
    """Let the pipeline generate below the 5 s floor the released checkpoint declares."""
    from diffusers.modular_pipelines.minimax_h3.modular_pipeline import MiniMaxH3ModularPipeline

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


def load_engine() -> str | None:
    """Load transformer + both VAEs at startup. Returns an error string, or None when the pipeline is up."""
    global PIPE, LOAD_ERROR, LOADED_IN

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

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

        _lower_duration_floor(MIN_DURATION)
        blocks = _generator_blocks()
        print(f"[h3] loading {[c.name for c in blocks.expected_components]} from {MODEL_REPO} ...", flush=True)

        pipe = blocks.init_pipeline(MODEL_REPO, components_manager=ComponentsManager(), collection="h3")
        # Both autoencoders carry `_keep_in_fp32_modules` and stay float32 - a bfloat16 audio VAE decodes the
        # soundtrack far too quiet.
        pipe.load_components(dtype=torch.bfloat16)
        pipe.transformer.set_attention_backend(ATTENTION)

        if PLACEMENT == "startup":
            # Only the transformer: `spaces` writes a second on-disk copy of every startup-resident CUDA tensor, and
            # packing all 77 GB busts the storage quota. The ~10 GB of fp32 VAEs move on the first request.
            pipe.transformer.to("cuda")

        PIPE = pipe
        LOADED_IN = time.time() - started
        print(f"[h3] 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: {type(error).__name__}: {error}"
        )
    return LOAD_ERROR


def status() -> str:
    """One line for the header pill."""
    if DEV_MODE:
        return "dev mode · no weights loaded"
    if LOAD_ERROR:
        return LOAD_ERROR
    if PIPE is None:
        return f"loading {MODEL_REPO} (transformer + VAEs, ~77 GB)"
    return (
        f"ready · bfloat16, unquantized · attention `{ATTENTION}` · GPU `{GPU_SIZE}` · "
        f"loaded in {LOADED_IN:.0f}s · conditioner `{CONDITIONER_SPACE}`"
    )


def is_ready() -> bool:
    return DEV_MODE or (PIPE is not None and LOAD_ERROR is None)


# ---------------------------------------------------------------------------
# Conditioner - the other half, over the gradio API
# ---------------------------------------------------------------------------
@cache
def _shared_conditioner():
    from gradio_client import Client

    return Client(CONDITIONER_SPACE)


def _conditioner(ip_token: str | None):
    """A conditioner client billed to the caller when we could read their token, to this Space's IP otherwise.

    Token forwarding through a header is what the ZeroGPU docs prescribe for a Space calling a Space; a per-request
    client is cheap next to the encode it books.
    """
    if not ip_token:
        return _shared_conditioner()
    from gradio_client import Client

    return Client(CONDITIONER_SPACE, headers={"x-ip-token": ip_token})


def _encode_key(prompt, first, last, canvas, num_frames) -> str:
    """Everything `/encode` reads. Seed, steps and the LoRA rack are deliberately absent: it sees none of them."""
    digest = hashlib.sha256()
    digest.update(prompt.strip().encode("utf-8"))
    for path in (first, last):
        digest.update(b"\0")
        if path:
            # The keyframes are already cover-cropped by the time they get here, so the bytes are stable.
            with open(path, "rb") as handle:
                while chunk := handle.read(1 << 20):
                    digest.update(chunk)
    digest.update(f"\0{canvas}\0{num_frames}".encode("utf-8"))
    return digest.hexdigest()


def encode(prompt, first, last, canvas, num_frames, ip_token=None):
    """`/encode`: a safetensors file with `prompt_embeds` + `text_token_tags` and the resolved geometry.

    Cached on its own inputs. With `rewrite_prompt=False` the conditioner is one forward pass of Qwen3-VL and
    nothing samples, so the same request has the same answer - and re-seeding or restacking LoRAs, which is most of
    what anyone does, stops costing the 40-odd seconds this call takes. A hit books no ZeroGPU time at all, on the
    caller's quota or ours. The embeds are a few megabytes each, so the cache is cheap to hold.
    """
    from gradio_client import handle_file
    from safetensors import safe_open

    key = _encode_key(prompt, first, last, canvas, num_frames)
    if key in _encode_cache:
        _encode_cache.move_to_end(key)
        print("[h3] conditioning served from cache", flush=True)
        return _encode_cache[key]

    path, plan = _conditioner(ip_token).predict(
        prompt=prompt,
        image_path=handle_file(first) if first else None,
        last_image_path=handle_file(last) if last else None,
        canvas=canvas,
        num_frames=num_frames,
        rewrite_prompt=False,
        api_name="/encode",
    )
    with safe_open(path, framework="pt") as handle:
        conditioning = (handle.get_tensor("prompt_embeds"), handle.get_tensor("text_token_tags"), handle.metadata(), plan)

    _encode_cache[key] = conditioning
    while len(_encode_cache) > ENCODE_CACHE_SIZE:
        _encode_cache.popitem(last=False)
    return conditioning


# ---------------------------------------------------------------------------
# Denoise
# ---------------------------------------------------------------------------
# Seconds of GPU one request needs, from the rows of the packed sequence it denoises: linear in the rows for the
# matmuls, quadratic for the attention. These are for an *eager* transformer - the published coefficients of this
# pipeline are fitted against an AOT-compiled block package, which this Space does not build, and booking against
# those would under-book by roughly threefold. `_calibrate` corrects them from what requests actually take.
_STEP_LINEAR, _STEP_QUADRATIC = 3.5e-4, 1.2e-8
_DECODE_BASE, _DECODE_PER_REFERENCE_CLIP = 15, 15
_REFERENCE_CLIP_PIXELS = 960 * 544 * 124
# An unfused adapter adds a pair of matmuls to every linear layer of every step.
_UNFUSED_STEP_COST = 0.15
# A LoRA that is not resident yet is read, converted and injected on booked time; folding one into the weights is
# one `B @ A` per target module. Both are one-offs that the rack stops paying once it settles.
_LORA_LOAD, _LORA_FOLD = 30, 25
_COLD_WORKER, _PAD = 12, 10
# The longest a single request may book. The full canvas at the longest clip estimates past this on an eager
# transformer, and asking for more than the pool allows fails the booking outright.
MAX_BOOKING = int(os.getenv("H3_MAX_BOOKING", "1500"))

# What the estimate is multiplied by, learned from the requests that ran. Starts neutral and is corrected on the
# first report; the clamp keeps one freak request from wrecking the booking for the next.
_calibration = 1.0
_CALIBRATION_WEIGHT, _CALIBRATION_RANGE = 0.4, (0.5, 4.0)


def _model_seconds(first, last, height, width, num_frames, steps, loras=()) -> float:
    """The denoise and decode this request is worth, before calibration and before the one-off LoRA costs."""
    height, width, num_frames, steps = int(height), int(width), int(num_frames), int(steps)
    patches = (height // 32) * (width // 32)
    rows = latent_frames(num_frames) * patches + (int(first is not None) + int(last is not None)) * patches
    denoise_seconds = steps * (_STEP_LINEAR * rows + _STEP_QUADRATIC * rows**2)
    denoise_seconds *= 1 + _UNFUSED_STEP_COST * len(loras)
    pixels = height * width * num_frames
    return denoise_seconds + _DECODE_BASE + _DECODE_PER_REFERENCE_CLIP * pixels / _REFERENCE_CLIP_PIXELS


def _calibrate(predicted: float, measured: float) -> None:
    """Pull the multiplier toward what the last request actually took."""
    global _calibration
    if predicted <= 0 or measured <= 0:
        return
    low, high = _CALIBRATION_RANGE
    ratio = min(max(measured / predicted, low), high)
    _calibration = min(max((1 - _CALIBRATION_WEIGHT) * _calibration + _CALIBRATION_WEIGHT * ratio, low), high)
    print(f"[h3] booking calibration now {_calibration:.2f} (predicted {predicted:.0f}s, took {measured:.0f}s)", flush=True)


def estimate_duration(prompt_embeds, text_token_tags, first, last, height, width, num_frames, steps, seed, loras=()):
    # Booked as if nothing were folded yet and as if this request is the one that folds: both are one-offs, and a
    # request cut off mid-denoise costs more than a slightly generous booking.
    one_offs = (_LORA_LOAD + _LORA_FOLD) * len(loras)
    booked = _model_seconds(first, last, height, width, num_frames, steps, loras) * _calibration + one_offs
    booked = max(60, int(booked) + _COLD_WORKER + _PAD)
    if booked > MAX_BOOKING:
        # An eager transformer at the top canvas and the longest clip asks for more than a booking is allowed to be.
        # Ask for the ceiling and say so: a request cut off at the ceiling is at least a legible failure.
        print(f"[h3] this request estimates {booked}s, over the {MAX_BOOKING}s ceiling - booking the ceiling", flush=True)
        booked = MAX_BOOKING
    return booked


@spaces.GPU(duration=estimate_duration, size=GPU_SIZE)
def denoise(prompt_embeds, text_token_tags, first, last, height, width, num_frames, steps, seed, loras=()):
    """The only thing on booked time: the LoRA stack, the denoise loop and the two decoders.

    Only the generated media comes back - a `@spaces.GPU` return crosses a process boundary by pickling, and the
    pipeline state still holds the packed latents and the rotary grid on the card.
    """
    import torch

    if PLACEMENT == "startup":
        PIPE.vae.to("cuda")
        PIPE.audio_vae.to("cuda")
    else:
        PIPE.to("cuda")

    applied = lora_stack.apply(PIPE, loras)

    # Timed separately from the rack work above, which is a one-off: this is the number the booking model is
    # calibrated against.
    model_started = time.time()
    state = PIPE(
        prompt_embeds=prompt_embeds.to("cuda"),
        text_token_tags=text_token_tags,
        image=first,
        last_image=last,
        height=int(height),
        width=int(width),
        num_frames=int(num_frames),
        num_inference_steps=int(steps),
        generator=torch.Generator("cpu").manual_seed(int(seed)),
    )
    model_seconds = time.time() - model_started
    return state.get("videos")[0], state.get("audio")[0].cpu(), state.get("sampling_rate"), applied, model_seconds


# ---------------------------------------------------------------------------
# One request
# ---------------------------------------------------------------------------
def generate(prompt, first_path, last_path, canvas, duration, steps, seed, loras=(), ip_token=None):
    """Conditioner -> denoise -> mux. Returns `(video_path, report, facts)`."""
    if LOAD_ERROR:
        raise RuntimeError(LOAD_ERROR)
    if PIPE is None and not DEV_MODE:
        raise RuntimeError("The denoiser is still loading - watch the Space logs.")
    if not prompt or not prompt.strip():
        raise ValueError("MiniMax-H3 always takes a prompt, keyframes or not.")

    canvas = canvas if canvas in CANVASES else DEFAULT_CANVAS
    if first_path:
        first_path, canvas = fit_keyframe(first_path, canvas)
    if last_path:
        last_path, canvas = fit_keyframe(last_path, canvas)
    num_frames = snap_frames(duration)

    if DEV_MODE:
        height, width = CANVASES[canvas]
        facts = {
            "seed": int(seed),
            "geometry": f"{width}x{height} · {num_frames} frames · {num_frames / FPS:.2f} s",
            "timing": "dev mode",
            "loras": ", ".join(f"{spec.label} @ {spec.weight:g}" for spec in loras) or "none",
        }
        return None, f"[dev] {prompt.strip()[:120]}", facts

    # Downloads happen here, in the main process: booked GPU time is not for pulling files off the network. A slot
    # whose file cannot be fetched is dropped with a note rather than taking the request down with it.
    resolved, notes = [], []
    for spec in loras:
        try:
            resolved.append(spec.fetch())
        except Exception as error:
            print(f"[lora] {spec.label} could not be fetched: {type(error).__name__}: {error}", flush=True)
            notes.append(f"{spec.label} unavailable ({type(error).__name__})")

    cached = _encode_key(prompt, first_path, last_path, canvas, num_frames) in _encode_cache
    encoded_at = time.time()
    prompt_embeds, text_token_tags, metadata, plan = encode(
        prompt, first_path, last_path, canvas, num_frames, ip_token=ip_token
    )
    encode_seconds = time.time() - encoded_at
    encode_note = "cached" if cached else f"{plan.get('num_text_tokens', '?')} text tokens"
    height, width, num_frames = (int(metadata[key]) for key in ("height", "width", "num_frames"))

    def keyframe(path):
        # The conditioner looked at exactly this image; the conditioning latents have to be of the same one.
        from PIL import Image, ImageOps

        return ImageOps.exif_transpose(Image.open(path)).convert("RGB") if path else None

    denoised_at = time.time()
    frames, audio, sampling_rate, applied, model_seconds = denoise(
        prompt_embeds,
        text_token_tags,
        keyframe(first_path),
        keyframe(last_path),
        height,
        width,
        num_frames,
        steps,
        seed,
        resolved,
    )
    denoise_seconds = time.time() - denoised_at
    applied = "; ".join([applied, *notes])
    # A folded rack no longer carries the per-step adapter cost, so the expectation it is measured against must not
    # carry it either - otherwise the multiplier drifts down and the next rack change under-books.
    unfused = () if "(folded)" in applied else resolved
    _calibrate(_model_seconds(first_path, last_path, height, width, num_frames, steps, unfused), model_seconds)

    from diffusers.utils import encode_video

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

    facts = {
        "seed": int(seed),
        "geometry": f"{width}x{height} · {num_frames} frames · {num_frames / FPS:.2f} s",
        "timing": f"encode {encode_seconds:.0f}s{' (cached)' if cached else ''} · denoise + decode {model_seconds:.0f}s",
        "loras": applied,
    }
    report = (
        f"{width}x{height} · {num_frames} frames ({num_frames / FPS:.2f} s) · {int(steps)} steps · "
        f"seed {int(seed)}\n"
        f"conditioner {encode_seconds:.0f}s ({encode_note}) · "
        f"denoise + decode {model_seconds:.0f}s ({model_seconds / max(1, int(steps)):.1f} s/step) · "
        f"rack {denoise_seconds - model_seconds:.0f}s\n"
        f"LoRA stack: {applied}"
    )
    print(f"[h3] {report}", flush=True)
    return path, report, facts