yn4989/minimax-h3

1"""MiniMax-H3, split deployment — denoising2"""3 4from __future__ import annotations5 6import os7import tempfile8import time9import traceback10 11# First, and at module level. `import spaces` patches `torch.cuda` before any GPU is attached, which is what lets the12# 72 GiB load happen at **startup** rather than on GPU time; it also has to precede anything that initializes CUDA.13import spaces14import gradio as gr15 16MODEL_REPO = os.environ.get("H3_MODEL_REPO", "diffusers-internal-dev/MiniMax-H3")17CONDITIONER_SPACE = os.environ.get("H3_CONDITIONER", "multimodalart/qwen3vl-conditioner")18# `lazy` moves all 72.16 GiB onto the card on the first GPU call and leaves it there; `offload` hands placement to19# `ComponentsManager.enable_auto_cpu_offload` instead. Neither puts anything on the card at *startup*, which is20# deliberate — see `load_models`: the 150 GB storage quota, not the 95 GiB card, is what rules that out here.21PLACEMENT = os.environ.get("H3_PLACEMENT", "pack").lower()22# cuDNN's fused attention is 10-20% faster than the SDPA default on this pool and needs nothing installed.23ATTENTION = os.environ.get("H3_ATTENTION", "_native_cudnn").lower()24GPU_DURATION = int(os.environ.get("H3_GPU_DURATION", "900"))25GPU_SIZE = os.environ.get("H3_GPU_SIZE", "xlarge")26ON_SPACES = bool(os.environ.get("SPACE_ID"))27 28CANVASES = {29    # 16:930    "960x544 · 16:9 fast": (544, 960),31    "1024x576 · 16:9 fast": (576, 1024),32    "1152x640 · 16:9": (640, 1152),33    "1280x704 · 16:9": (704, 1280),34    "1344x768 · 16:9 full": (768, 1344),35    # 9:1636    "544x960 · 9:16 fast": (960, 544),37    "640x1152 · 9:16": (1152, 640),38    "768x1344 · 9:16 full": (1344, 768),39    # 1:140    "544x544 · 1:1 fast": (544, 544),41    "768x768 · 1:1 full": (768, 768),42    # 4:3 / 3:443    "768x576 · 4:3 fast": (576, 768),44    "1024x768 · 4:3 full": (768, 1024),45    "576x768 · 3:4 fast": (768, 576),46    "768x1024 · 3:4 full": (1024, 768),47    # 21:948    "1152x512 · 21:9 fast": (512, 1152),49    "1536x672 · 21:9 full": (672, 1536),50}51DEFAULT_CANVAS = "960x544 · 16:9 fast"52FPS, FRAMES_PER_CHUNK, LATENTS_PER_CHUNK = 24, 17, 553MAX_UI_DURATION = 1454 55 56def snap_frames(seconds: float) -> int:57    """The frame count MiniMax-H3's video VAE can decode: the next `17 * n + 5` at 24 fps."""58    frames = max(1, round(float(seconds) * FPS))59    while frames % FRAMES_PER_CHUNK != LATENTS_PER_CHUNK:60        frames += 161    return frames62 63 64PIPE = None65MANAGER = None66LOAD_ERROR: str | None = None67LOADED_IN: float | None = None68CLIENT = None69 70 71def status() -> str:72    if LOAD_ERROR:73        return LOAD_ERROR74    if PIPE is None:75        return f"Loading `{MODEL_REPO}` (transformer + VAEs, 77.3 GB). Watch the Space logs."76    import h3_aoti77 78    return (79        f"Ready · transformer + VAEs **bfloat16, unquantized** · placement `{PLACEMENT}` · attention `{ATTENTION}` · "80        f"{h3_aoti.status()} · loaded in {LOADED_IN:.0f}s · conditioner `{CONDITIONER_SPACE}`"81    )82 83 84def load_models() -> str | None:85    """Load the denoising half. At **startup**, but *not* onto the card.86 87    `MiniMaxH3GeneratorBlocks` declares `transformer`, `vae`, `audio_vae`, `scheduler`, `audio_scheduler` and88    `video_processor`, so `load_components` fetches exactly those subfolders out of the shared89    `modular_model_index.json` — `text_encoder/` and `transformer_ref/` are never touched.90 91    Both autoencoders carry `_keep_in_fp32_modules` over every module, so the `dtype` below is refused for them and92    they stay float32: a bfloat16 audio VAE decodes the soundtrack roughly 20 dB too quiet.93 94    Nothing is moved onto the card here, which is the one place this Space departs from the ZeroGPU idiom, and the95    reason is storage rather than memory. `spaces`' startup `torch.pack()` writes every startup-resident CUDA tensor96    to a **second copy on disk** and only deletes the downloaded originals afterwards; 77.3 GB of weights plus a97    77.3 GB pack is 154.6 GB against a 150 GB quota, and the Space is evicted mid-pack with `OSError: [Errno 28] No98    space left on device` out of `os.posix_fallocate`. Deleting the shards first does not help either: the pack's own99    cleanup walks the still-open mappings and `lstat`s them, so an unlinked blob turns into `FileNotFoundError:100    ... (deleted)`. Placement therefore happens on the first GPU call, where it costs about 10 s of PCIe and then101    persists across every later request in the same worker.102    """103    global PIPE, MANAGER, LOAD_ERROR, LOADED_IN104 105    if PIPE is not None or LOAD_ERROR is not None:106        return LOAD_ERROR107 108    token = os.environ.get("HF_TOKEN")109    if not token:110        LOAD_ERROR = f"**`HF_TOKEN` secret is missing** and `{MODEL_REPO}` is private. Add it and restart."111        return LOAD_ERROR112 113    started = time.time()114    try:115        import torch116        from diffusers import ComponentsManager117 118        from h3_split_blocks import MiniMaxH3GeneratorBlocks119 120        manager = ComponentsManager()121        blocks = MiniMaxH3GeneratorBlocks()122        print(f"[gen] loading {[c.name for c in blocks.expected_components]} from {MODEL_REPO} ...", flush=True)123        pipe = blocks.init_pipeline(MODEL_REPO, components_manager=manager, collection="h3")124        pipe.load_components(dtype=torch.bfloat16, token=token)125        pipe.transformer.set_attention_backend(ATTENTION)126 127        # Still startup, still free: an AoTI package carries no weights and opens its compiled archive lazily inside128        # the GPU worker, so pointing the 50-block stack at it is CPU work. Off unless `H3_AOTI=1`.129        import h3_aoti130 131        h3_aoti.maybe_load(pipe.transformer)132 133        if PLACEMENT == "pack":134            # Idiomatic ZeroGPU startup placement, scoped to the transformer only. `spaces` packs every135            # startup-resident CUDA tensor into a second on-disk copy; packing all 77.3 GB (transformer + fp32136            # VAEs) busts the 150 GB storage quota (77.3 + 77.3 + shards), but the 61.7 GB transformer alone137            # packs to ~123 GB total and fits. The VAEs (~10 GB) take the lazy path on first GPU call, ~2 s.138            # With AoTI the packed transformer pairs with the precompiled blocks: no placement, no compile,139            # first request runs at steady state.140            pipe.transformer.to("cuda")141 142        if PLACEMENT == "offload":143            manager.enable_auto_cpu_offload(device="cuda")144            _arm_decode_hooks(pipe)145 146        PIPE, MANAGER = pipe, manager147        LOADED_IN = time.time() - started148        print(f"[gen] ready in {LOADED_IN:.0f}s", flush=True)149    except Exception as error:150        traceback.print_exc()151        LOAD_ERROR = f"**Loading `{MODEL_REPO}` failed** after {time.time() - started:.0f}s: `{type(error).__name__}: {error}`"152    return LOAD_ERROR153 154 155def _arm_decode_hooks(pipe):156    """Make the offload hooks fire for the two VAEs.157 158    `enable_auto_cpu_offload` installs accelerate hooks, which wrap `forward`. The decode blocks call159    `components.vae.decode(...)` and `components.audio_vae.decode(...)` directly, so the hook never runs and the VAE160    is still on the host when the latents arrive on the card.161    """162    for name in ("vae", "audio_vae"):163        module = getattr(pipe, name)164        inner = module.decode165 166        def armed(*args, _module=module, _decode=inner, **kwargs):167            hook = getattr(_module, "_hf_hook", None)168            if hook is not None:169                hook.pre_forward(_module)170            return _decode(*args, **kwargs)171 172        module.decode = armed173 174 175def conditioner():176    """The other half, over the gradio API. Cached — building a `Client` costs a round trip to the Space config."""177    global CLIENT178    if CLIENT is None:179        from gradio_client import Client180 181        CLIENT = Client(CONDITIONER_SPACE)  # public Space, no org token: the request runs on the caller side quota182    return CLIENT183 184 185def encode_remote(prompt, image_path, last_image_path, canvas, num_frames):186    """Ask the conditioner Space for `prompt_embeds` + `text_token_tags`. Off this Space's GPU time entirely."""187    from gradio_client import handle_file188    from safetensors import safe_open189 190    path, plan = conditioner().predict(191        prompt=prompt,192        image_path=handle_file(image_path) if image_path else None,193        last_image_path=handle_file(last_image_path) if last_image_path else None,194        canvas=canvas,195        num_frames=num_frames,196        api_name="/encode",197    )198    with safe_open(path, framework="pt") as handle:199        metadata = handle.metadata()200        return handle.get_tensor("prompt_embeds"), handle.get_tensor("text_token_tags"), metadata, plan201 202 203 204 205# Fitted on live probes (5 configs spanning canvas, duration and steps; max residual 3.7 s):206# gpu_seconds = A + B * steps * tokens + C * steps * tokens^2, where tokens is the packed video row count.207# PLACEMENT_ALLOWANCE covers the one-time 72 GiB lazy .to("cuda") a cold worker pays inside its first call.208_DUR_A, _DUR_B, _DUR_C = -6.023, 2.0877e-4, 2.1221e-9209_PLACEMENT_ALLOWANCE, _PAD = 12, 10  # pack mode: only the ~10 GB VAEs move on a cold worker210 211 212def get_duration(prompt_embeds, text_token_tags, image, last_image, height, width, num_frames, steps, seed, *a, **k):213    latent_frames = (int(num_frames) - 5) // 17 * 5 + 2214    patches = (int(height) // 32) * (int(width) // 32)215    tokens = latent_frames * patches216    tokens += (int(image is not None) + int(last_image is not None)) * patches217    st = int(steps) * tokens218    compute = _DUR_A + _DUR_B * st + _DUR_C * st * tokens219    return max(60, int(compute) + _PLACEMENT_ALLOWANCE + _PAD)220 221 222@spaces.GPU(duration=get_duration, size=GPU_SIZE)223def _generate(prompt_embeds, text_token_tags, image, last_image, height, width, num_frames, steps, seed):224    """The only thing on GPU time: the packed-sequence denoise loop and the two decoders.225 226    Only the three generated outputs come back. A `@spaces.GPU` return crosses a process boundary by pickling, and227    the full `PipelineState` still holds the packed latents, the rotary grid and the row indices on the card.228    """229    import torch230 231    if PLACEMENT == "lazy":232        # 72.16 GiB across PCIe on the first request of a worker, a no-op walk on every one after it.233        PIPE.to("cuda")234    elif PLACEMENT == "pack":235        # Transformer was packed at startup; only the ~10 GB of fp32 VAEs walk across on a cold worker.236        PIPE.vae.to("cuda")237        PIPE.audio_vae.to("cuda")238 239    state = PIPE(240        prompt_embeds=prompt_embeds.to("cuda"),241        text_token_tags=text_token_tags,242        image=image,243        last_image=last_image,244        height=height,245        width=width,246        num_frames=num_frames,247        num_inference_steps=int(steps),248        generator=torch.Generator("cpu").manual_seed(int(seed)),249    )250    return state.get("videos")[0], state.get("audio")[0].cpu(), state.get("sampling_rate")251 252 253def generate(prompt, image_path=None, last_image_path=None, canvas=DEFAULT_CANVAS, duration=5, steps=28, seed=42, progress=gr.Progress(track_tqdm=True)):254    if LOAD_ERROR:255        raise gr.Error(LOAD_ERROR)256    if PIPE is None:257        raise gr.Error("The denoiser is still loading.")258    if not prompt or not prompt.strip():259        raise gr.Error("MiniMax-H3 always takes a prompt, keyframes or not.")260 261    from PIL import Image262 263    from diffusers.utils import encode_video264 265    num_frames = snap_frames(duration)266 267    progress(0.0, desc=f"Conditioning on {CONDITIONER_SPACE} ...")268    conditioned = time.time()269    prompt_embeds, text_token_tags, metadata, plan = encode_remote(270        prompt, image_path, last_image_path, canvas, num_frames271    )272    condition_seconds = time.time() - conditioned273    height, width, num_frames = (int(metadata[key]) for key in ("height", "width", "num_frames"))274 275    progress(0.1, desc=f"Denoising {steps} steps at {width}x{height}, {num_frames} frames ...")276    started = time.time()277    frames, audio, sampling_rate = _generate(278        prompt_embeds,279        text_token_tags,280        Image.open(image_path) if image_path else None,281        Image.open(last_image_path) if last_image_path else None,282        height,283        width,284        num_frames,285        steps,286        seed,287    )288    generate_seconds = time.time() - started289 290    directory = os.path.join(tempfile.gettempdir(), "h3-outputs")291    os.makedirs(directory, exist_ok=True)292    path = os.path.join(directory, f"h3-{int(time.time() * 1000)}.mp4")293    encode_video(frames, fps=FPS, output_path=path, audio=audio, audio_sample_rate=sampling_rate)294 295    report = (296        f"`{width}x{height}`, {num_frames} frames ({num_frames / FPS:.3f} s), {int(steps)} steps · "297        f"conditioner {condition_seconds:.0f}s ({plan['num_text_tokens']} tokens) · "298        f"denoise + decode {generate_seconds:.0f}s ({generate_seconds / int(steps):.1f} s/step) · seed {int(seed)}"299    )300    print(f"[gen] {report}", flush=True)301    return path, report302 303 304 305def _fit_keyframe(image_path, current_canvas):306    """Cover-crop an uploaded keyframe to the closest supported aspect ratio and select that ratio's307    smallest (fastest) canvas, unless the user already picked a matching ratio."""308    if not image_path:309        return gr.update(), gr.update()310    from PIL import Image as _Image311 312    img = _Image.open(image_path)313    aspect = img.width / img.height314    fastest = {}315    for label, (h, w) in CANVASES.items():316        r = w / h317        if r not in fastest or w * h < fastest[r][1][0] * fastest[r][1][1]:318            fastest[r] = (label, (h, w))319    ratio = min(fastest, key=lambda r: abs(r - aspect))320    label, (h, w) = fastest[ratio]321 322    cur_h, cur_w = CANVASES[current_canvas]323    if abs(cur_w / cur_h - aspect) <= abs(ratio - aspect):324        label = current_canvas325        h, w = cur_h, cur_w326 327    target = w / h328    if abs(img.width / img.height - target) <= 1e-3:329        return gr.update(), gr.update(value=label)330    if True:331        if img.width / img.height > target:332            new_w = int(img.height * target)333            left = (img.width - new_w) // 2334            img = img.crop((left, 0, left + new_w, img.height))335        else:336            new_h = int(img.width / target)337            top = (img.height - new_h) // 2338            img = img.crop((0, top, img.width, top + new_h))339        img.save(image_path)340    return gr.update(value=image_path), gr.update(value=label)341 342 343load_models()344 345INTRO = """# MiniMax-H3346 347<div align="center">348  <a href="https://huggingface.co/MiniMaxAI/MiniMax-H3"><strong>[ model ]</strong></a> &nbsp;349  <a href="PAPER_URL_PLACEHOLDER"><strong>[ paper ]</strong></a> &nbsp;350  <a href="https://www.minimax.io"><strong>[ project ]</strong></a>351</div>352 353**MiniMax-H3** is a 33B parameter state of the art video generation model that produces video and a354fully synchronized soundtrack (ambience, foley, speech).355"""356 357CSS = """358.main.fillable {max-width: 1250px !important}359.dark .gradio-container { color: var(--body-text-color); }360"""361 362with gr.Blocks(title="MiniMax-H3") as demo:363    gr.Markdown(INTRO)364 365    with gr.Row():366        with gr.Column():367            prompt = gr.Textbox(368                label="Prompt",369                lines=3,370                value="A red fox trotting through a snowy pine forest at dawn, snow crunching underfoot",371            )372            with gr.Row():373                image = gr.Image(label="First frame (optional)", type="filepath")374                last_image = gr.Image(label="Last frame (optional)", type="filepath")375            run = gr.Button("Generate", variant="primary")376            with gr.Accordion("Advanced options", open=False):377                canvas = gr.Dropdown(label="Canvas", choices=list(CANVASES), value=DEFAULT_CANVAS)378                duration = gr.Slider(label="Duration (s)", minimum=2, maximum=MAX_UI_DURATION, step=1, value=5)379                steps = gr.Slider(label="Steps", minimum=10, maximum=40, step=1, value=28)380                seed = gr.Number(label="Seed", value=42, precision=0)381            382        with gr.Column():383            video = gr.Video(label="Video + soundtrack")384            report = gr.Markdown(visible=False)385 386    image.upload(_fit_keyframe, [image, canvas], [image, canvas])387 388    gr.Examples(389        examples=[390            ["A red fox trotting through a snowy pine forest at dawn, snow crunching underfoot", None, None, "1344x768 · 16:9 full"],391            ["A busy night market, neon signs reflecting in puddles, sizzling street food", None, None, "768x1344 · 9:16 full"],392            ["A cellist playing a slow melody in an empty concert hall", None, None, "768x768 · 1:1 full"],393            ["The fox looks around, then trots deeper into the forest", "examples/first.png", None, "1344x768 · 16:9 full"],394            ["A slow seamless camera move from the first view to the last", "examples/first.png", "examples/last.png", "1344x768 · 16:9 full"],395        ],396        inputs=[prompt, image, last_image, canvas],397        outputs=[video, report],398        fn=generate,399        cache_examples=True,400        cache_mode="lazy",401    )402 403    run.click(404        generate,405        [prompt, image, last_image, canvas, duration, steps, seed],406        [video, report],407        api_name="generate",408    )409 410 411if __name__ == "__main__":412    demo.launch(show_error=True, theme=gr.themes.Citrus(), css=CSS)413