hao-ai-lab-FastVideo/examples/inference/basic/basic_minimax_h3_t2v.py

# SPDX-License-Identifier: Apache-2.0
"""Generate video and audio from text with MiniMax H3."""

from __future__ import annotations

import argparse
import os
from pathlib import Path

from fastvideo import VideoGenerator
from fastvideo.api import (
    CompileConfig,
    EngineConfig,
    GenerationRequest,
    GeneratorConfig,
    OffloadConfig,
    OutputConfig,
    ParallelismConfig,
    PipelineSelection,
    SamplingConfig,
)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--model-path", default="MiniMaxAI/MiniMax-H3")
    # Rank-reduced AdaLN checkpoint (-39% params, -23 GiB VRAM): pass
    #   --model-path noctuashap/MiniMax-H3-pruned-r16
    # (or a local dir produced by
    #   scripts/checkpoint_conversion/convert_minimax_h3_adaln_rank.py).
    # adaln_rank is read from the checkpoint config; no other flags needed.
    # Rank-reduced checkpoints are inference-only: training needs the
    # full-rank release.
    parser.add_argument("--prompt", required=True)
    parser.add_argument("--output", default="outputs/minimax_h3_t2v")
    parser.add_argument("--height", type=int, default=768)
    parser.add_argument("--width", type=int, default=1344)
    parser.add_argument("--num-frames", type=int, default=124)
    parser.add_argument("--steps", type=int, default=50)
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument("--num-gpus", type=int, default=4)
    parser.add_argument(
        "--execution-backend",
        choices=("mp", "uni", "ray"),
        default=None,
        help="mp for one node; ray for a Ray cluster (two DGX Sparks). "
        "Default: ray when RAY_ADDRESS is set, otherwise mp",
    )
    parser.add_argument("--torch-compile", action="store_true", help="torch.compile the DiT transformer path")
    parser.add_argument("--compile-vae",
                        action=argparse.BooleanOptionalAction,
                        default=True,
                        help="compile the video VAE decoder independently of the DiT (on by default; "
                        "the Spark lazy-load path needs this registered before first materialize)")
    parser.add_argument("--compile-mode",
                        default=None,
                        help='torch.compile mode, e.g. "reduce-overhead" for CUDA graphs')
    parser.add_argument("--inference-torch-compile",
                        action="store_true",
                        help="regional fullgraph torch.compile of each DiT block after load (the #1718 "
                        "training-port semantics: no kwargs; fullgraph + emulate_precision_casts injected). "
                        "First generation pays the inductor JIT (~1-2 min); use --repeats >= 2 and time "
                        "the last repeat. FASTVIDEO_INFERENCE_TORCH_COMPILE=1 is equivalent")
    parser.add_argument("--lazy-module-load",
                        action=argparse.BooleanOptionalAction,
                        default=None,
                        help="load each heavy component on first use and free it after the last stage that "
                        "needs it, so peak memory is the largest overlapping set instead of the sum of every "
                        "component. Omit for auto (on for unified-memory devices such as GB10; off on discrete "
                        "GPUs). Costs a reload per generation; pass --no-lazy-module-load to keep every "
                        "component resident")
    parser.add_argument("--repeats",
                        type=int,
                        default=1,
                        help="generate N times; with --torch-compile the first run pays "
                        "compilation, so steady-state is the last repeat")
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    output_dir = Path(args.output)
    output_dir.mkdir(parents=True, exist_ok=True)

    # Boot-time run configuration folded into FastVideoArgs (the same
    # experimental-dict route basic_fasth3.py uses for the VSA knobs).
    experimental: dict[str, object] = {}
    if args.inference_torch_compile:
        experimental["inference_torch_compile"] = True

    execution_backend = args.execution_backend or ("ray" if os.environ.get("RAY_ADDRESS") else "mp")
    generator = VideoGenerator.from_config(
        GeneratorConfig(
            model_path=args.model_path,
            pipeline=PipelineSelection(experimental=experimental),
            engine=EngineConfig(
                num_gpus=args.num_gpus,
                execution_backend=execution_backend,
                use_fsdp_inference=args.num_gpus > 1,
                parallelism=ParallelismConfig(tp_size=1, sp_size=args.num_gpus),
                offload=OffloadConfig(
                    dit=False,
                    dit_layerwise=False,
                    text_encoder=True,
                    vae=True,
                    pin_cpu_memory=False,
                    lazy_module_load=args.lazy_module_load,
                ),
                compile=CompileConfig(
                    enabled=args.torch_compile,
                    mode=args.compile_mode,
                    vae_enabled=args.compile_vae,
                ),
            ),
        ))
    try:
        request = GenerationRequest(
            prompt=args.prompt,
            negative_prompt="",
            sampling=SamplingConfig(
                height=args.height,
                width=args.width,
                num_frames=args.num_frames,
                fps=24,
                num_inference_steps=args.steps,
                guidance_scale=1.0,
                batch_cfg=False,
                seed=args.seed,
            ),
            output=OutputConfig(
                output_path=str(output_dir / "minimax_h3_t2v.mp4"),
                save_video=True,
                return_frames=False,
            ),
        )
        result = generator.generate(request)
        print(f"Output written to: {result.video_path}")
        if result.generation_time is not None:
            # machine-readable: benchmark harnesses parse this line to separate
            # generation from model-load time (last occurrence = steady state)
            print(f"Generation time: {result.generation_time:.2f}s")
        # Worker-lifetime allocator high-water mark: torch never resets it
        # between generations, so it includes model load and every prior run.
        if result.peak_memory_mb is not None:
            print(f"Peak memory: {result.peak_memory_mb:.1f} MB")
        for _ in range(args.repeats - 1):
            result = generator.generate(request)
            if result.generation_time is not None:
                print(f"Generation time: {result.generation_time:.2f}s")
            # Same worker-lifetime peak as above; it does not reset per repeat.
            if result.peak_memory_mb is not None:
                print(f"Peak memory: {result.peak_memory_mb:.1f} MB")
    finally:
        generator.shutdown()


if __name__ == "__main__":
    main()