Aryan Kumar
361f919c88
Keep the four-step uniform AdaLN cache, and use explicit DMD rungs only when the snapshot declares its own scheduler shifts.
237 lines
10 KiB
Python
237 lines
10 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""End-to-end MiniMax-H3 (FastH3) generation with the Apple Silicon MLX runtime.
|
|
|
|
Accepts a text prompt and produces an MP4 with H.264 video at 24 fps and
|
|
stereo AAC audio at 32 kHz. One heavyweight model phase is resident at a time.
|
|
|
|
python examples/inference/basic/mlx_fasth3.py \
|
|
--model-root ~/models/FastH3-Preview-v0.2 \
|
|
--mlx-checkpoint ~/models/FastH3-MLX/int8 \
|
|
--prompt '(S1) A red panda says <d>[English] Fast H3 is amazing.</d>' \
|
|
--height 480 --width 832 --num-frames 124 --seed 2026 \
|
|
--output-path ~/fasth3_outputs/int8.mp4
|
|
|
|
Conditioning uses the streamed Qwen3-VL text encoder on first use and caches
|
|
the resulting embeddings under --prompt-cache-dir for instant reuse.
|
|
|
|
``--fast`` is temporal fast mode. It keeps full-duration audio while
|
|
denoising fewer video frames, then uses MLX RIFE 4.25 to reconstruct the
|
|
requested frame count. A 1280x720 request runs on H3's 1280x736 grid and is
|
|
center-cropped after decode.
|
|
|
|
``--fast-spatial`` is spatial fast mode, ``--fast``'s spatial twin. It
|
|
denoises and decodes on the smallest 32px-aligned canvas covering
|
|
height/width divided by ``--fast-spatial-scale``, then resamples the decoded
|
|
frames up to the requested size in pixel space. The two modes compose.
|
|
This trades fine detail for speed: the output carries the reduced canvas's
|
|
detail budget and reads softer than a native-resolution render, so it stays
|
|
off by default.
|
|
|
|
This entrypoint currently supports text-to-video-with-audio only. It does not
|
|
yet wire FL2VA, Ref2VA, or two-pass refinement. The four-step preview stays
|
|
the default. FastH3 8-Step V2 uses ``mlx_fasth3_8step.py``.
|
|
|
|
VSA is off by default; existing dense MLX checkpoints remain supported.
|
|
H3 uses fused MLX RMSNorm, which can change BF16 rounding compared with the
|
|
older explicit normalization path. Convert with ``--include-vsa`` and pass
|
|
``--vsa`` to enable the sparse path. Attention activations stay BF16; INT6/INT8/INT4 apply only to
|
|
linear weights, including the optional gate projection.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
from pathlib import Path
|
|
|
|
|
|
def _dense_layers(value: str) -> tuple[int, ...]:
|
|
layers = tuple(int(part.strip()) for part in value.split(",") if part.strip())
|
|
if any(layer < 0 for layer in layers):
|
|
raise argparse.ArgumentTypeError("dense layer indices must be non-negative")
|
|
return layers
|
|
|
|
|
|
def build_parser() -> argparse.ArgumentParser:
|
|
parser = argparse.ArgumentParser(description=__doc__,
|
|
formatter_class=argparse.RawDescriptionHelpFormatter)
|
|
parser.add_argument("--model-root", type=Path, default=Path.home() / "models/FastH3-Preview-v0.2",
|
|
help="H3 snapshot root (vae/, audio_vae/, text_encoder/, tokenizer/)")
|
|
parser.add_argument("--mlx-checkpoint", type=Path, required=True,
|
|
help="pre-quantized MLX DiT directory (int8/int6/int4 mlx_h3_dit format)")
|
|
parser.add_argument(
|
|
"--prompt",
|
|
required=True,
|
|
help="H3 text prompt; use (S1) and <d>[Language] words</d> for explicit dialogue",
|
|
)
|
|
parser.add_argument("--output-path", type=Path, required=True)
|
|
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=0)
|
|
parser.add_argument("--steps", type=int, default=4, help="denoise steps (trained ladder = 4)")
|
|
parser.add_argument(
|
|
"--fast",
|
|
action=argparse.BooleanOptionalAction,
|
|
default=False,
|
|
help="denoise fewer video frames, then use MLX RIFE to restore the target frame count; audio stays full length",
|
|
)
|
|
parser.add_argument("--fast-factor", type=int, default=2,
|
|
help="temporal reduction target for --fast (default: 2)")
|
|
parser.add_argument("--fast-sharpen", type=float, default=0.6,
|
|
help="unsharp strength after RIFE interpolation (0 disables)")
|
|
parser.add_argument(
|
|
"--fast-spatial",
|
|
action=argparse.BooleanOptionalAction,
|
|
default=False,
|
|
help="denoise and decode at height/width // fast-spatial-scale on H3's 32px grid, "
|
|
"then resample the decoded frames up to the requested size; composes with --fast. "
|
|
"Trades fine detail for speed",
|
|
)
|
|
parser.add_argument("--fast-spatial-scale", type=int, default=2,
|
|
help="spatial reduction factor for --fast-spatial (default: 2)")
|
|
parser.add_argument("--fast-spatial-upsample-mode",
|
|
choices=("lanczos", "cubic", "bilinear", "nearest"),
|
|
default="lanczos",
|
|
help="pixel interpolation kernel for the post-decode upsample")
|
|
parser.add_argument("--fast-spatial-sharpen", type=float, default=0.4,
|
|
help="unsharp strength after the upsample (0 disables)")
|
|
parser.add_argument("--rife-weights-dir", type=Path, default=None,
|
|
help="optional local mlx-community/RIFE-4.25 snapshot")
|
|
parser.add_argument("--vae-dtype", choices=("fp32", "fp16", "bf16"), default="fp32")
|
|
parser.add_argument("--video-decode-backend", choices=("h3-vae", "taeh3"), default="h3-vae",
|
|
help="full H3 VAE or approximate TAEH3 preview decoder; audio is unchanged")
|
|
parser.add_argument("--taeh3-checkpoint", type=Path, default=None,
|
|
help="local TAEH3 safetensors; otherwise download hash-verified upstream weights")
|
|
parser.add_argument("--taeh3-chunk-size", type=int, default=5,
|
|
help="latent frames per TAEH3 chunk; causal memory persists across chunks")
|
|
parser.add_argument("--prompt-cache-dir", type=Path, default=None,
|
|
help="directory for reusable prompt embedding caches")
|
|
parser.add_argument(
|
|
"--tiled-video-decode",
|
|
action=argparse.BooleanOptionalAction,
|
|
default=True,
|
|
help="decode with the reference 256px overlapping VAE tiles (disable only for diagnostics)",
|
|
)
|
|
parser.add_argument(
|
|
"--vsa",
|
|
action=argparse.BooleanOptionalAction,
|
|
default=False,
|
|
help=(
|
|
"enable MiniMax H3 VSA; requires a VSA-capable MLX checkpoint "
|
|
"from --include-vsa"
|
|
),
|
|
)
|
|
parser.add_argument(
|
|
"--vsa-sparsity",
|
|
type=float,
|
|
default=0.9,
|
|
help="VSA sparsity in [0, 1); 0.9 is the trained FastH3 policy",
|
|
)
|
|
parser.add_argument(
|
|
"--vsa-tile-size",
|
|
type=int,
|
|
default=64,
|
|
choices=(64, 256),
|
|
help="VSA tile size in tokens",
|
|
)
|
|
parser.add_argument(
|
|
"--vsa-prefix-mode",
|
|
choices=("exempt", "compete"),
|
|
default="exempt",
|
|
help=(
|
|
"prefix-key policy: always keep (exempt) or FLOP-matched top-k "
|
|
"(compete)"
|
|
),
|
|
)
|
|
parser.add_argument(
|
|
"--vsa-dense-first-n-steps",
|
|
type=int,
|
|
default=0,
|
|
help="run the first N denoise steps dense",
|
|
)
|
|
parser.add_argument(
|
|
"--vsa-dense-layers",
|
|
type=_dense_layers,
|
|
default=(),
|
|
help="comma-separated layer indices forced dense",
|
|
)
|
|
parser.add_argument(
|
|
"--vsa-impl",
|
|
choices=("auto", "reference", "simd"),
|
|
default="auto",
|
|
help=(
|
|
"sparse attention implementation; auto uses chunked gather+SDPA "
|
|
"(simd is opt-in)"
|
|
),
|
|
)
|
|
parser.add_argument("--vae-tile-height", type=int, default=256,
|
|
help="overlapping video-VAE tile height in pixels")
|
|
parser.add_argument("--vae-tile-width", type=int, default=256,
|
|
help="overlapping video-VAE tile width in pixels")
|
|
parser.add_argument("--inter-step-cooldown-s", type=float, default=0.0,
|
|
help="optional pause between denoise steps")
|
|
return parser
|
|
|
|
|
|
def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
|
|
return build_parser().parse_args(argv)
|
|
|
|
|
|
def run(args: argparse.Namespace) -> None:
|
|
from fastvideo.mlx_runtime.minimax_h3_pipeline import MiniMaxH3MLXPipeline
|
|
|
|
pipeline = MiniMaxH3MLXPipeline(
|
|
model_root=args.model_root,
|
|
mlx_dit_checkpoint=args.mlx_checkpoint,
|
|
vae_dtype=args.vae_dtype,
|
|
video_decode_backend=args.video_decode_backend,
|
|
taeh3_checkpoint=args.taeh3_checkpoint,
|
|
taeh3_chunk_size=args.taeh3_chunk_size,
|
|
prompt_cache_dir=args.prompt_cache_dir,
|
|
)
|
|
result = pipeline.generate(
|
|
args.prompt,
|
|
output_path=args.output_path,
|
|
height=args.height,
|
|
width=args.width,
|
|
num_frames=args.num_frames,
|
|
seed=args.seed,
|
|
num_steps=args.steps,
|
|
tiled_video_decode=args.tiled_video_decode,
|
|
vae_tile_height=args.vae_tile_height,
|
|
vae_tile_width=args.vae_tile_width,
|
|
inter_step_cooldown_s=args.inter_step_cooldown_s,
|
|
fast=args.fast,
|
|
fast_factor=args.fast_factor,
|
|
fast_sharpen=args.fast_sharpen,
|
|
rife_weights_dir=args.rife_weights_dir,
|
|
fast_spatial=args.fast_spatial,
|
|
fast_spatial_scale=args.fast_spatial_scale,
|
|
fast_spatial_upsample_mode=args.fast_spatial_upsample_mode,
|
|
fast_spatial_sharpen=args.fast_spatial_sharpen,
|
|
vsa=args.vsa,
|
|
vsa_sparsity=args.vsa_sparsity,
|
|
vsa_tile_size=args.vsa_tile_size,
|
|
vsa_prefix_mode=args.vsa_prefix_mode,
|
|
vsa_dense_first_n_steps=args.vsa_dense_first_n_steps,
|
|
vsa_dense_layers=args.vsa_dense_layers,
|
|
vsa_impl=args.vsa_impl,
|
|
)
|
|
print(json.dumps({
|
|
"video_path": result.video_path,
|
|
"timings_s": {k: round(v, 2) for k, v in result.timings.items()},
|
|
"peak_memory_gib": {k: round(v, 2) for k, v in result.peak_memory_gib.items()},
|
|
"vsa": result.vsa,
|
|
"video_decode_backend": result.video_decode_backend,
|
|
"audio_samples": int(result.waveform.shape[-1]),
|
|
}, indent=2))
|
|
|
|
|
|
def main() -> None:
|
|
run(parse_args())
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|