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

# SPDX-License-Identifier: Apache-2.0
"""Reference-conditioned FastH3 OmniRef (Ref2VA) generation on Apple Silicon MLX.

A FastH3 OmniRef PDD export carries only what distillation changed:
``transformer_ref/``, the two scheduler configs and ``fastvideo_inference.json``.
The text encoder, tokenizer, processor and both VAEs are base MiniMax-H3's, so
this script composes the two into one directory of symlinks - the same shape the
CUDA example builds - and then runs the MLX runtime over it.

There is no FP4 on Apple Silicon, so the DiT converts to INT8 (or INT6/INT4),
the same formats the text-to-video-audio MLX path uses. The DiT has to be in
the packed MLX layout first:

    python scripts/checkpoint_conversion/convert_minimax_h3_mlx.py \\
      --model-root <composed>/transformer_ref --out ~/models/FastH3-OmniRef-MLX \\
      --formats int8 --include-vsa

    python examples/inference/basic/mlx_fasth3_omniref.py \\
      --export ./FastVideo-FastH3-Omni-8-Step-V1 --composed-dir ./FastH3-OmniRef-composed \\
      --base-model-path ./MiniMax-H3 --mlx-checkpoint ~/models/FastH3-OmniRef-MLX/int8 \\
      --image reference.png --audio voice.wav \\
      --prompt "a person on a sofa, lit by a warm lamp, speaking to camera" \\
      --output-path ~/fasth3_outputs/omniref.mp4

References are ordered: pass them in the order the model should read them, with
``--image`` / ``--video`` / ``--audio`` (for example ``--video dance.mp4
--image outfit.png --audio voice.wav``). At least one image or video reference
is required; audio references only accompany them.

Add ``--compose-only`` to build the composed directory (the conversion above
needs it) and stop, with no checkpoint, prompt or reference.

The PDD ladder comes from the export's contract, so ``--steps`` is optional:
without it the trained fused-block count runs.
"""

from __future__ import annotations

import argparse
import json
from pathlib import Path
from typing import Any

CONTRACT = "fastvideo_inference.json"
MANIFESTS = ("modular_model_index.json", "model_index.json")
EXPORT_COMPONENTS = ("transformer_ref", "scheduler", "audio_scheduler")
BASE_COMPONENTS = ("text_encoder", "tokenizer", "processor", "vae", "audio_vae")


class _AppendReference(argparse.Action):
    """Collect --image/--video/--audio into one list that keeps command-line order."""

    def __call__(self, parser, namespace, value, option_string=None):
        references = list(getattr(namespace, self.dest, None) or [])
        references.append((self.const, value))
        setattr(namespace, self.dest, references)


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--export", type=Path,
                        help="FastH3 OmniRef PDD export directory; composes it with --base-model-path")
    parser.add_argument("--base-model-path", type=Path,
                        help="base MiniMax-H3 snapshot (vae/, audio_vae/, text_encoder/, tokenizer/)")
    parser.add_argument("--composed-dir", type=Path, default=Path("outputs/fasth3_omniref_mlx_model"),
                        help="where the composed directory of symlinks is written")
    parser.add_argument("--model-root", type=Path,
                        help="an already-composed directory; skips composing")
    parser.add_argument("--compose-only", action="store_true",
                        help="compose --export with --base-model-path into --composed-dir and exit")
    parser.add_argument("--mlx-checkpoint", type=Path,
                        help="packed MLX checkpoint converted from the composed transformer_ref/")
    for flag in ("image", "video", "audio"):
        parser.add_argument(f"--{flag}", dest="references", action=_AppendReference, const=flag, metavar="PATH",
                            help=f"an ordered {flag} reference (repeatable)")
    parser.add_argument("--prompt")
    parser.add_argument("--output-path", type=Path, default=Path("outputs/mlx_fasth3_omniref.mp4"))
    parser.add_argument("--height", type=int, default=480)
    parser.add_argument("--width", type=int, default=832)
    parser.add_argument("--num-frames", type=int, default=124)
    parser.add_argument("--seed", type=int, default=2026)
    parser.add_argument("--steps", type=int, default=0,
                        help="transformer forwards; 0 uses the export contract's trained count")
    parser.add_argument("--prompt-cache-dir", type=Path)
    parser.add_argument("--metal-wired-limit-gib", type=float,
                        help="keep this much unified memory wired so macOS does not page the DiT out mid-denoise")
    parser.add_argument("--mlx-cache-limit-gib", type=float,
                        help="bound MLX's freed-buffer cache; unbounded by default")
    parser.add_argument("--quant-gemm", choices=("auto", "dequant", "qmm"), default=None,
                        help="INT8/INT6/INT4 linear path: auto (measured default), dequant (dequantize + dense "
                        "GEMM for wide inputs), qmm (mx.quantized_matmul). Unset defers to "
                        "FASTVIDEO_MLX_QUANT_GEMM. Ignored by a bf16 checkpoint.")
    parser.add_argument("--reference-max-pixels", default=None,
                        help="opt-in cap on each image/video reference canvas (down only): a pixel count, or "
                        "'match' for the clip's own area. Faster, but references lose detail versus the trained "
                        "canvas.")
    parser.add_argument("--vae-dtype", default="fp32")
    parser.add_argument("--vae-decoder", type=Path,
                        help="standalone video decoder (FastH3 decoder[-nvfp4].safetensors or a LightVAE decoder) "
                        "replacing the base vae/ decoder")
    parser.add_argument("--vae-encoder", type=Path,
                        help="standalone video encoder (LightVAE) replacing the base vae/ encoder; with both flags "
                        "vae/ only needs config.json")
    parser.add_argument("--vae-decoder-quantization", choices=("nvfp4", "mxfp8", "int8", "int4"),
                        help="quantize the decoder block linears of a dense decoder file as it loads")
    parser.add_argument("--vsa", action=argparse.BooleanOptionalAction, default=None,
                        help="VSA sparse attention (needs a checkpoint converted with --include-vsa). "
                        "On by default when the contract is a VSA contract; --no-vsa runs dense.")
    parser.add_argument("--vsa-sparsity", type=float, default=None,
                        help="generated-video sparsity; default is the contract's trained value")
    parser.add_argument("--vsa-tile-size", type=int, default=None,
                        help="VSA tile: 64, 128 or 256. Must equal the contract's trained tile; "
                        "default is the contract's.")
    parser.add_argument("--vsa-ref-keep-rate", type=float, default=None,
                        help="fraction of each reference video's tiles a query keeps (p2_multi_region); "
                        "default is the contract's trained value")
    parser.add_argument("--vsa-impl", choices=("auto", "reference", "simd", "sdpa", "nax"), default="auto",
                        help="sparse attention backend; auto uses the SIMD Metal kernel for 128-token tiles and "
                        "gather+SDPA otherwise, reference forces gather+SDPA; sdpa shares K/V with a tile mask "
                        "and nax uses M5 TensorOps over selected blocks (both opt-in, incompatible with --reuse-prefix-kv)")
    parser.add_argument("--reuse-prefix-kv", action="store_true",
                        help="opt-in, lossy: record the text/reference rows' per-block K/V on step 1 and run later "
                        "steps on the target audio/video rows only (faster, slightly different output)")
    parser.add_argument("--prefix-kv-refresh-every", type=int, default=0,
                        help="re-record the prefix K/V with a full forward every N steps (0 never)")
    parser.add_argument("--prefix-kv-cache-max-gib", type=float, default=4.0,
                        help="refuse prefix K/V reuse when the estimated cache exceeds this many GiB")
    parser.add_argument("--prefix-kv-storage", choices=("bf16", "fp16", "int8"), default="bf16",
                        help="prefix K/V cache storage dtype (int8 halves the cache)")
    parser.add_argument("--inter-step-cooldown-s", type=float, default=0.0)
    return parser


def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
    parser = build_parser()
    args = parser.parse_args(argv)
    if args.compose_only:
        if args.export is None or args.base_model_path is None:
            parser.error("--compose-only needs --export and --base-model-path")
        return args
    if args.mlx_checkpoint is None or args.prompt is None:
        parser.error("--mlx-checkpoint and --prompt are required unless --compose-only is set")
    if not args.references:
        parser.error("pass at least one ordered reference with --image, --video, or --audio")
    if not any(kind != "audio" for kind, _ in args.references):
        parser.error("Ref2VA needs at least one image or video reference")
    if args.model_root is None and (args.export is None or args.base_model_path is None):
        parser.error("give --model-root, or --export together with --base-model-path")
    if args.vsa_impl in ("sdpa", "nax") and args.reuse_prefix_kv:
        parser.error(f"--vsa-impl {args.vsa_impl} does not support --reuse-prefix-kv")
    return args


def _link(destination: Path, source: Path) -> None:
    if not source.exists():
        raise FileNotFoundError(f"Missing checkpoint component: {source}")
    source = source.resolve()
    if destination.is_symlink():
        if destination.resolve() == source:
            return
        destination.unlink()
    elif destination.exists():
        raise FileExistsError(f"{destination} exists and is not a symlink; choose another --composed-dir.")
    destination.symlink_to(source, target_is_directory=source.is_dir())


def compose_model_dir(export_dir: Path, base_dir: Path, composed_dir: Path) -> Path:
    """One MiniMax-H3 model directory: distilled components from the export, the rest from the base."""
    contract = export_dir / CONTRACT
    if not contract.is_file():
        raise FileNotFoundError(f"{export_dir} has no {CONTRACT}; it is not a FastH3 distilled export.")
    composed_dir.mkdir(parents=True, exist_ok=True)
    for name in (*EXPORT_COMPONENTS, CONTRACT):
        _link(composed_dir / name, export_dir / name)
    for name in MANIFESTS:
        for root in (export_dir, base_dir):
            if (root / name).is_file():
                _link(composed_dir / name, root / name)
                break
    for name in BASE_COMPONENTS:
        _link(composed_dir / name, base_dir / name)
    return composed_dir


def load_contract(export_dir: Path) -> dict[str, Any]:
    """The export's PDD contract, validated as a Ref2VA one."""
    path = export_dir / CONTRACT
    if not path.is_file():
        raise FileNotFoundError(f"{export_dir} has no {CONTRACT}")
    contract = json.loads(path.read_text(encoding="utf-8"))
    if contract.get("model_type") != "ref2va":
        raise ValueError(f"{path} is not a Ref2VA contract; use the example matching the checkpoint.")
    return contract


def resolve_reference_max_pixels(value: str | None, width: int, height: int) -> int | None:
    """``--reference-max-pixels``: unset keeps the trained canvas; 'match' uses the clip's area."""
    if value is None:
        return None
    if value == "match":
        return width * height
    pixels = int(value)
    if pixels <= 0:
        raise ValueError(f"--reference-max-pixels must be positive or 'match', got {value!r}.")
    return pixels


def resolve_model_root(args: argparse.Namespace) -> Path:
    if args.model_root is not None:
        return args.model_root
    return compose_model_dir(args.export, args.base_model_path, args.composed_dir)


def run(args: argparse.Namespace) -> None:
    if args.compose_only:
        print(f"Composed model directory: {resolve_model_root(args)}")
        return
    from fastvideo.mlx_runtime.minimax_h3_pipeline import MiniMaxH3MLXPipeline
    from fastvideo.mlx_runtime.minimax_h3_reference import MiniMaxH3Reference

    model_root = resolve_model_root(args)
    export_dir = args.export if args.export is not None else model_root
    contract = load_contract(export_dir)
    if args.export is not None:
        print(f"Composed model directory: {model_root}")
        print(f"Contract: {contract['num_inference_steps']} fused blocks of a {contract['pdd_steps']}-interval grid, "
              f"{contract['attention_backend']} with sparsity {contract['vsa_sparsity']}, "
              f"{contract['vsa_tile_size']}-token tiles")
    references = [MiniMaxH3Reference(source=path, media_type=kind) for kind, path in args.references]
    args.output_path.parent.mkdir(parents=True, exist_ok=True)

    pipeline = MiniMaxH3MLXPipeline(
        model_root=model_root,
        mlx_dit_checkpoint=args.mlx_checkpoint,
        vae_dtype=args.vae_dtype,
        vae_decoder_path=args.vae_decoder,
        vae_encoder_path=args.vae_encoder,
        vae_decoder_quantization=args.vae_decoder_quantization,
        prompt_cache_dir=args.prompt_cache_dir,
        metal_wired_limit_gib=args.metal_wired_limit_gib,
        mlx_cache_limit_gib=args.mlx_cache_limit_gib,
        quant_gemm=args.quant_gemm,
    )
    try:
        kwargs: dict[str, Any] = {}
        if args.steps:
            kwargs["num_steps"] = args.steps
        elif "transformer_forwards" in contract:
            # The contract's fused-block count is the ladder the student was distilled on.
            kwargs["num_steps"] = int(contract["transformer_forwards"])
        result = pipeline.generate(
            args.prompt,
            output_path=args.output_path,
            references=references,
            height=args.height,
            width=args.width,
            num_frames=args.num_frames,
            seed=args.seed,
            vsa=args.vsa,
            vsa_sparsity=args.vsa_sparsity,
            vsa_tile_size=args.vsa_tile_size,
            vsa_ref_keep_rate=args.vsa_ref_keep_rate,
            vsa_impl=args.vsa_impl,
            inter_step_cooldown_s=args.inter_step_cooldown_s,
            reuse_prefix_kv=args.reuse_prefix_kv,
            prefix_kv_refresh_every=args.prefix_kv_refresh_every,
            prefix_kv_cache_max_gib=args.prefix_kv_cache_max_gib,
            prefix_kv_storage=args.prefix_kv_storage,
            reference_max_pixels=resolve_reference_max_pixels(args.reference_max_pixels, args.width, args.height),
            **kwargs,
        )
    finally:
        pipeline.close()
    print(json.dumps({
        "video_path": result.video_path,
        "references": [kind for kind, _ in args.references],
        "timings_s": {key: round(value, 2) for key, value in result.timings.items()},
        "peak_memory_gib": {key: round(value, 2) for key, value in result.peak_memory_gib.items()},
        "vsa": result.vsa,
        "prefix_kv": result.prefix_kv,
        "audio_samples": int(result.waveform.shape[-1]),
    }, indent=2))


def main() -> None:
    run(parse_args())


if __name__ == "__main__":
    main()