376e445f0e
Co-authored-by: Aryan Kumar <aryan5v@users.noreply.github.com> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
534 lines
26 KiB
Python
534 lines
26 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""Eight-forward Ref2VA video+audio generation with a FastH3 OmniRef PDD student.
|
|
|
|
A FastH3 OmniRef PDD export is a Parallel Decoding Distillation (PDD) student of
|
|
MiniMax-H3's reference-conditioned ``transformer_ref`` partition. Its directory
|
|
(local, or a Hugging Face repo) carries only what distillation changed:
|
|
``transformer_ref/`` (output heads widened to the fine grid, trained VSA
|
|
compression gates), the two scheduler configs, and ``fastvideo_inference.json``.
|
|
The text encoder, tokenizer, processor, and both VAEs are base MiniMax-H3's.
|
|
|
|
This script composes the two into one local model directory of symlinks and
|
|
runs the recipe the contract records: fused blocks of the fine grid (one
|
|
transformer forward each), the video/audio shifts, and VIDEO_SPARSE_ATTN_H3
|
|
with its sparsity and tile size, where every reference video is its own
|
|
sparse region. The contract's 128-token tiles run only on the sm_100a/sm_103a
|
|
CUDA kernel (B200/B300/GB200/GB300) of a fastvideo-kernel build with the
|
|
Blackwell VSA extension; the VSA-H3 backend raises at the first attention call
|
|
when either is missing.
|
|
With ``--num-gpus`` above 1 the DiT is sharded across the GPUs (FSDP) and the
|
|
sequence is split across them (sequence parallelism).
|
|
|
|
References are ordered. Pass them in order with --image / --video / --audio,
|
|
for example ``--video dance.mp4 --image outfit.png --audio voice.wav``.
|
|
|
|
``--profile {spark,rtx-pro-6000,rtx-5090,rtx-4090}`` runs the consumer stack on
|
|
one of those devices: a packed NVFP4 export, FP4 sparse attention with the dense
|
|
prefix rows split off at 8-bit Q/K, and the pre-packed NVFP4 decoder, with the
|
|
device's own offload policy and geometry. The 4090, having no FP4 hardware, runs
|
|
the closest equivalent from a dense/FP8 export (FP8 linears, bf16 VSA, the dense decoder). See
|
|
docs/inference/fasth3_omniref.md.
|
|
|
|
``--lossless-accel`` turns on every bit-identical speed-up the pipeline has
|
|
(LOSSLESS_ACCEL_ENV below, plus resident encoders, pinned host memory and the
|
|
sequence-parallel VAE); the output is bitwise the same as without it. It keeps
|
|
the 63 GB Qwen3-VL text encoder resident on every GPU, so it needs about
|
|
125 GB per GPU at four GPUs (B200/B300 class).
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import os
|
|
from collections.abc import Sequence
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
from fastvideo import VideoGenerator
|
|
from fastvideo.api import (
|
|
CompileConfig,
|
|
ComponentConfig,
|
|
EngineConfig,
|
|
GenerationRequest,
|
|
GeneratorConfig,
|
|
InputConfig,
|
|
OffloadConfig,
|
|
OutputConfig,
|
|
ParallelismConfig,
|
|
PipelineSelection,
|
|
QuantizationConfig,
|
|
SamplingConfig,
|
|
)
|
|
from fastvideo.configs.pipelines.minimax_h3 import parse_base_model_revision
|
|
from fastvideo.pipelines.basic.minimax_h3 import MiniMaxH3Reference
|
|
|
|
CONTRACT = "fastvideo_inference.json"
|
|
DEFAULT_HEIGHT, DEFAULT_WIDTH = 480, 832
|
|
# Components the distilled export replaces, and the ones it shares with base MiniMax-H3.
|
|
EXPORT_COMPONENTS = ("transformer_ref", "scheduler", "audio_scheduler")
|
|
BASE_COMPONENTS = ("text_encoder", "tokenizer", "processor", "vae", "audio_vae")
|
|
MANIFESTS = ("modular_model_index.json", "model_index.json")
|
|
# Components the composed directory's manifest must declare, besides its transformer_ref.
|
|
MANIFEST_COMPONENTS = ("scheduler", "audio_scheduler", *BASE_COMPONENTS)
|
|
# Bit-identical speed-ups that FastVideo reads from the environment. Workers
|
|
# inherit it, so these must be set before the generator starts.
|
|
LOSSLESS_ACCEL_ENV = {
|
|
# Fused NVLink Ulysses all-to-all instead of NCCL plus permute copies.
|
|
"FASTVIDEO_ULYSSES_A2A": "auto",
|
|
# RoPE, AdaLN modulation / gated residual and SwiGLU kernels that round to
|
|
# bf16 exactly where the eager ops do (not the FP32 round-once fusions).
|
|
"FASTVIDEO_MINIMAX_H3_EXACT_KERNELS": "all",
|
|
# VSA tiles scattered straight into the kernels' heads-first layout.
|
|
"FASTVIDEO_H3_VSA_HEADS_FIRST_TILE": "1",
|
|
# VAE spatial tiles split across the sequence-parallel ranks.
|
|
"FASTVIDEO_H3_VAE_TILE_PARALLEL": "1",
|
|
# Content-keyed reuse of repeat reference encodes.
|
|
"FASTVIDEO_H3_REF2VA_MEMO_ENTRIES": "16",
|
|
}
|
|
|
|
|
|
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("--model-path",
|
|
required=True,
|
|
help="FastH3 OmniRef PDD export: a local directory or a Hugging Face repo id")
|
|
parser.add_argument("--revision", default=None, help="revision of --model-path when it is a Hugging Face repo")
|
|
parser.add_argument("--base-model-path",
|
|
default=None,
|
|
help="base MiniMax-H3 snapshot (local directory or repo id). Default: the repo and revision "
|
|
"pinned by the export's base_model_revision")
|
|
parser.add_argument("--base-revision",
|
|
default=None,
|
|
help="revision of --base-model-path when it is a repo id (the export's pinned revision "
|
|
"applies only to its own base repo); without --base-model-path, overrides that pin")
|
|
parser.add_argument("--composed-dir",
|
|
default=None,
|
|
help="where to write the composed model directory of symlinks (default: "
|
|
"<output>/fasth3_omniref_model)")
|
|
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", required=True)
|
|
parser.add_argument("--output", default="outputs/fasth3_omniref_pdd")
|
|
parser.add_argument("--height", type=int, default=None)
|
|
parser.add_argument("--width", type=int, default=None)
|
|
parser.add_argument("--num-frames", type=int, default=124)
|
|
parser.add_argument("--seed", type=int, default=0)
|
|
parser.add_argument("--num-gpus", type=int, default=1, help="sequence-parallel GPUs (must divide 56 heads)")
|
|
parser.add_argument("--lossless-accel",
|
|
action="store_true",
|
|
help="enable every bit-identical speed-up (output unchanged; about 125 GB per GPU)")
|
|
parser.add_argument("--profile",
|
|
choices=sorted(DEVICE_PROFILES),
|
|
default=None,
|
|
help="run the full NVFP4 consumer stack for one device, with its offload policy and "
|
|
"environment (see docs/inference/fasth3_omniref.md); the checkpoint must be "
|
|
"the packed NVFP4 export on Blackwell, or a dense/FP8 export on the 4090")
|
|
parser.add_argument("--attention-prefix",
|
|
choices=("default", "fp8", "bf16"),
|
|
default="default",
|
|
help="numerics for the dense attention prefix rows (text, reference and audio). "
|
|
"'fp8' is the validated-quality mode and the device-profile default; 'bf16' is "
|
|
"higher precision but slower; 'default' keeps the profile's choice")
|
|
parser.add_argument("--decoder",
|
|
default=None,
|
|
help="pre-packed NVFP4 VAE decoder (decoder-nvfp4.safetensors) used by --profile "
|
|
"instead of the VAE folder's own decoder")
|
|
parser.add_argument("--adaln-table",
|
|
default=None,
|
|
help="precomputed AdaLN modulation table, so the 26 GB of adaln_proj never loads "
|
|
"(write one once with FASTVIDEO_H3_ADALN_DUMP). Strongly recommended: the 24/32 GB "
|
|
"profiles need it to keep the DiT resident")
|
|
return parser
|
|
|
|
|
|
def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
|
|
parser = build_parser()
|
|
args = parser.parse_args(argv)
|
|
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")
|
|
return args
|
|
|
|
|
|
def _snapshot(path_or_repo: str, revision: str | None, allow_patterns: list[str]) -> Path:
|
|
local = Path(path_or_repo).expanduser()
|
|
if local.is_dir():
|
|
return local
|
|
from huggingface_hub import snapshot_download
|
|
|
|
return Path(snapshot_download(repo_id=path_or_repo, revision=revision, allow_patterns=allow_patterns))
|
|
|
|
|
|
def load_contract(export_dir: Path) -> dict[str, Any]:
|
|
path = export_dir / CONTRACT
|
|
if not path.is_file():
|
|
raise FileNotFoundError(f"{export_dir} has no {CONTRACT}; it is not a FastH3 distilled export.")
|
|
contract = json.loads(path.read_text(encoding="utf-8"))
|
|
if "pdd_steps" not in contract or contract.get("model_type") != "ref2va":
|
|
raise ValueError(f"{path} is not a Ref2VA PDD contract; use the example matching the checkpoint.")
|
|
return contract
|
|
|
|
|
|
def base_model_source(args: argparse.Namespace, contract: dict[str, Any]) -> tuple[str, str | None]:
|
|
"""The base snapshot and revision: --base-model-path as given, else the export's pin."""
|
|
base_repo, base_revision = base_model_from_contract(contract)
|
|
if args.base_model_path:
|
|
return args.base_model_path, args.base_revision
|
|
return base_repo, args.base_revision or base_revision
|
|
|
|
|
|
def base_model_from_contract(contract: dict[str, Any]) -> tuple[str, str]:
|
|
"""The base repo id and revision that the export pins as ``hf://<repo id>@<revision>``."""
|
|
return parse_base_model_revision(contract.get("base_model_revision"))
|
|
|
|
|
|
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 _declares_components(manifest: Path) -> bool:
|
|
"""Whether a Diffusers manifest declares every component of the composed directory."""
|
|
declared = {
|
|
name
|
|
for name, spec in json.loads(manifest.read_text(encoding="utf-8")).items()
|
|
if isinstance(spec, list) and spec and spec[0] is not None
|
|
}
|
|
return set(MANIFEST_COMPONENTS) <= declared and bool({"transformer", "transformer_ref"} & declared)
|
|
|
|
|
|
def select_manifest(export_dir: Path, base_dir: Path) -> Path:
|
|
"""The export's manifest when it declares every composed component, else the base's."""
|
|
for root in (export_dir, base_dir):
|
|
for name in MANIFESTS:
|
|
if (root / name).is_file() and _declares_components(root / name):
|
|
return root / name
|
|
raise FileNotFoundError("Neither the export nor the base snapshot has a Diffusers model manifest that declares "
|
|
f"{', '.join(MANIFEST_COMPONENTS)} and transformer_ref.")
|
|
|
|
|
|
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."""
|
|
composed_dir.mkdir(parents=True, exist_ok=True)
|
|
for name in (*EXPORT_COMPONENTS, CONTRACT):
|
|
_link(composed_dir / name, export_dir / name)
|
|
manifest_source = select_manifest(export_dir, base_dir)
|
|
for name in MANIFESTS:
|
|
stale = composed_dir / name
|
|
if name != manifest_source.name and stale.is_symlink():
|
|
stale.unlink()
|
|
_link(composed_dir / manifest_source.name, manifest_source)
|
|
for name in BASE_COMPONENTS:
|
|
_link(composed_dir / name, base_dir / name)
|
|
return composed_dir
|
|
|
|
|
|
def resolve_model(args: argparse.Namespace) -> tuple[Path, dict[str, Any]]:
|
|
export_dir = _snapshot(args.model_path, args.revision,
|
|
[CONTRACT, *MANIFESTS, *(f"{name}/**" for name in EXPORT_COMPONENTS)])
|
|
contract = load_contract(export_dir)
|
|
if args.profile == "rtx-4090" and (export_dir / "transformer_ref" / "nvfp4_weights.safetensors").is_file():
|
|
raise ValueError("rtx-4090 requires a dense or FP8 OmniRef export. The packed NVFP4 export omits "
|
|
"the weights needed for FP8 conversion; pass the dense/FP8 export with --model-path.")
|
|
base_source, base_revision = base_model_source(args, contract)
|
|
base_dir = _snapshot(base_source, base_revision, [*MANIFESTS, *(f"{name}/**" for name in BASE_COMPONENTS)])
|
|
composed_dir = Path(args.composed_dir) if args.composed_dir else Path(args.output) / "fasth3_omniref_model"
|
|
return compose_model_dir(export_dir, base_dir, composed_dir), contract
|
|
|
|
|
|
def apply_lossless_accel_env() -> None:
|
|
"""Set LOSSLESS_ACCEL_ENV, keeping any value already in the environment."""
|
|
for name, value in LOSSLESS_ACCEL_ENV.items():
|
|
os.environ.setdefault(name, value)
|
|
|
|
|
|
_CONSUMER_ENV = {
|
|
# The checkpoint's trained attention backend; the contract fixes the tile size.
|
|
"FASTVIDEO_ATTENTION_BACKEND": "VIDEO_SPARSE_ATTN_H3",
|
|
# FP4 sparse attention on the consumer-Blackwell kernel, with the prefix split.
|
|
"FASTVIDEO_H3_VSA_FP4": "1",
|
|
"FASTVIDEO_H3_VSA_FP4_PREFIX": "fp8",
|
|
# QAD trained the video rows with two-level P and K smoothing; match it.
|
|
"FASTVIDEO_H3_VSA_FP4_TWO_LEVEL_P": "1",
|
|
"FASTVIDEO_H3_VSA_FP4_K_SMOOTH": "1",
|
|
# Bit-identical gate chunking; caps the 1.2 GiB transient at 768p10.
|
|
"FASTVIDEO_H3_GATE_CHUNK_ROWS": "8192",
|
|
# Bit-exact kernels, and none of the data-center-only routes.
|
|
"FASTVIDEO_MINIMAX_H3_FUSIONS": "all",
|
|
"FASTVIDEO_VSA_TRITON": "1",
|
|
"FASTVIDEO_VSA_PTX_BLACKWELL": "0",
|
|
"FASTVIDEO_FA4": "0",
|
|
}
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class DeviceProfile:
|
|
"""Everything that differs between consumer devices for one OmniRef clip."""
|
|
|
|
label: str
|
|
width: int
|
|
height: int
|
|
offload: OffloadConfig
|
|
env: dict[str, str]
|
|
# Defer the denoise modules until the sampler needs them, for the cards that
|
|
# cannot hold the text encoder and the DiT at once (the 24/32 GB recipes use it).
|
|
sequential_load: bool = False
|
|
# Linear quantization. sm_89 has no FP4 hardware, so the 4090 profile is the
|
|
# closest equivalent rather than the FP4 stack, in the DiT and - because the
|
|
# NVFP4 decoder's GEMMs need it too - in the VAE.
|
|
transformer_quant: str = "NVFP4"
|
|
layer_profile: str | None = "h3_dit_vsa"
|
|
nvfp4_decoder: bool = True
|
|
|
|
|
|
DEVICE_PROFILES: dict[str, DeviceProfile] = {
|
|
"spark": DeviceProfile(
|
|
label="NVIDIA DGX Spark (GB10, 128 GB unified)",
|
|
width=832,
|
|
height=480,
|
|
# 128 GB unified holds the encoder, the DiT and the decoder at once:
|
|
# the same resident policy as the released Spark recipe.
|
|
offload=OffloadConfig(dit=False,
|
|
dit_layerwise=False,
|
|
text_encoder=False,
|
|
vae=False,
|
|
pin_cpu_memory=False,
|
|
lazy_module_load=False),
|
|
env={
|
|
**_CONSUMER_ENV,
|
|
# Byte-identical at both shapes and both decoders; the reference
|
|
# memo saves 5-16% of warm wall time on repeated references.
|
|
**LOSSLESS_ACCEL_ENV,
|
|
# cuDNN mm_fp4 is 2.7-3x faster than CUTLASS on the H3 FFN shapes on GB10.
|
|
"FASTVIDEO_NVFP4_MM_BACKEND": "cudnn",
|
|
"FASTVIDEO_NVFP4_MM_ROW_SPLIT": "table",
|
|
"FASTVIDEO_H3_VAE_TILE_BATCH": "12",
|
|
"FASTVIDEO_H3_VAE_INT8_OVERLAY": "0",
|
|
},
|
|
),
|
|
"rtx-pro-6000": DeviceProfile(
|
|
label="NVIDIA RTX PRO 6000 (96 GB)",
|
|
width=1344,
|
|
height=768,
|
|
# 96 GB fits the encoder, the DiT and the decoder at once.
|
|
offload=OffloadConfig(dit=False,
|
|
dit_layerwise=False,
|
|
text_encoder=False,
|
|
vae=False,
|
|
pin_cpu_memory=True,
|
|
lazy_module_load=False),
|
|
env={**_CONSUMER_ENV, "FASTVIDEO_H3_VAE_TILE_BATCH": "8"},
|
|
),
|
|
"rtx-5090": DeviceProfile(
|
|
label="NVIDIA RTX 5090 (32 GB)",
|
|
width=832,
|
|
height=480,
|
|
# The encoder is not resident, and the denoise modules load only when the
|
|
# sampler needs them. The DiT itself stays whole: no layerwise streaming.
|
|
offload=OffloadConfig(dit=False, dit_layerwise=False, text_encoder=True, vae=False, pin_cpu_memory=True),
|
|
sequential_load=True,
|
|
env={
|
|
**_CONSUMER_ENV,
|
|
# Bit-identical at both shapes and both decoders on the 5090; the
|
|
# reference memo and VAE paths save 3-5 s on repeated requests.
|
|
**LOSSLESS_ACCEL_ENV,
|
|
"FASTVIDEO_H3_VAE_TILE_BATCH": "4",
|
|
},
|
|
),
|
|
"rtx-4090": DeviceProfile(
|
|
label="NVIDIA RTX 4090 (24 GB, sm_89, no FP4)",
|
|
width=832,
|
|
height=480,
|
|
# The 47 GB FP8 DiT cannot be resident on 24 GB: stream it, keeping the
|
|
# validated 16 blocks resident. The encoder stays non-resident.
|
|
offload=OffloadConfig(dit=False, dit_layerwise=True, text_encoder=True, vae=False, pin_cpu_memory=True),
|
|
sequential_load=True,
|
|
env={
|
|
# No FP4 hardware: FP8 linears and bf16 VSA attention instead, and the
|
|
# dense light VAE because the NVFP4 decoder's GEMMs need FP4 too.
|
|
"FASTVIDEO_ATTENTION_BACKEND": "VIDEO_SPARSE_ATTN_H3",
|
|
"FASTVIDEO_H3_VSA_FP4": "0",
|
|
"FASTVIDEO_H3_VSA_FINE_ATTENTION": "emulator-bf16",
|
|
"FASTVIDEO_MINIMAX_H3_FUSIONS": "all",
|
|
"FASTVIDEO_VSA_TRITON": "1",
|
|
"FASTVIDEO_VSA_PTX_BLACKWELL": "0",
|
|
"FASTVIDEO_FA4": "0",
|
|
# Validated 24 GB layerwise configuration: stream the packed weights
|
|
# too, keep 16 blocks resident, and chunk the FFN.
|
|
"FASTVIDEO_LAYERWISE_OFFLOAD_BUFFERS": "1",
|
|
"FASTVIDEO_LAYERWISE_RESIDENT_BLOCKS": "16",
|
|
"FASTVIDEO_H3_FFN_CHUNK_TOKENS": "16384",
|
|
"FASTVIDEO_H3_PINNED_SWAP": "0",
|
|
# A 24 GB card fragments easily once the DiT is streamed.
|
|
"PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True",
|
|
"FASTVIDEO_H3_VAE_TILE_BATCH": "4",
|
|
},
|
|
transformer_quant="FP8",
|
|
layer_profile=None,
|
|
nvfp4_decoder=False,
|
|
),
|
|
}
|
|
|
|
|
|
def apply_attention_prefix(choice: str) -> None:
|
|
"""Apply ``--attention-prefix``; ``default`` keeps the profile's choice."""
|
|
if choice != "default":
|
|
os.environ["FASTVIDEO_H3_VSA_FP4_PREFIX"] = choice
|
|
|
|
|
|
def apply_profile(profile: DeviceProfile, decoder: str | None, adaln_table: str | None = None) -> None:
|
|
"""Set a profile's environment, keeping any value already in the environment."""
|
|
for name, value in profile.env.items():
|
|
os.environ.setdefault(name, value)
|
|
if decoder and profile.nvfp4_decoder:
|
|
os.environ.setdefault("FASTVIDEO_H3_VAE_NVFP4_DECODER", decoder)
|
|
elif decoder:
|
|
print(f"--decoder ignored on {profile.label}: the NVFP4 decoder needs FP4 hardware; "
|
|
"this profile decodes with the VAE folder's own (dense) decoder")
|
|
if adaln_table:
|
|
# The AdaLN projections are a third of the DiT's bytes on disk; the
|
|
# precomputed table drops them entirely. Every profile takes it: on the
|
|
# 24/32 GB cards it is what makes residency possible, and elsewhere it
|
|
# buys headroom for the reference tokens.
|
|
os.environ.setdefault("FASTVIDEO_H3_ADALN_TABLE", adaln_table)
|
|
|
|
|
|
def build_generator_config(model_dir: Path,
|
|
num_gpus: int,
|
|
lossless_accel: bool = False,
|
|
profile: DeviceProfile | None = None) -> GeneratorConfig:
|
|
quantization = None
|
|
if profile is not None:
|
|
# `h3_dit_vsa` matches the packed export the consumer weights ship; the
|
|
# export's metadata names its bf16 blocks, so no skip list is needed.
|
|
quantization = QuantizationConfig(transformer_quant=profile.transformer_quant,
|
|
layer_profile=profile.layer_profile)
|
|
offload = profile.offload
|
|
# Explicit: on unified memory the auto rule would defer the DiT and load
|
|
# 50+ GB at the first denoise call even when everything fits.
|
|
experimental = {"h3_sequential_load": profile.sequential_load}
|
|
if num_gpus > 1:
|
|
experimental.update({"vae_parallel_decode": True, "vae_parallel_encode": True})
|
|
elif lossless_accel:
|
|
# Resident encoders skip streaming the 63 GB text encoder (and the
|
|
# VAEs) from host memory on every clip; pinned buffers copy the
|
|
# decoded clip at full PCIe speed; the parallel VAE paths reassemble
|
|
# the serial result exactly.
|
|
offload = OffloadConfig(dit=False, dit_layerwise=False, text_encoder=False, vae=False, pin_cpu_memory=True)
|
|
experimental = {"vae_parallel_decode": num_gpus > 1, "vae_parallel_encode": num_gpus > 1}
|
|
else:
|
|
offload = OffloadConfig(dit=False, dit_layerwise=False, text_encoder=True, vae=True, pin_cpu_memory=False)
|
|
experimental = {}
|
|
return GeneratorConfig(
|
|
model_path=str(model_dir),
|
|
engine=EngineConfig(
|
|
num_gpus=num_gpus,
|
|
# Multi-GPU device profiles run on Ray (the Spark pair spans two hosts).
|
|
execution_backend="ray" if num_gpus > 1 and profile is not None else "mp",
|
|
# A packed NVFP4 export is replicated across GPUs; the tokens are split
|
|
# by sequence parallelism. Data-center checkpoints keep mp and FSDP.
|
|
use_fsdp_inference=num_gpus > 1 and profile is None,
|
|
parallelism=ParallelismConfig(tp_size=1, sp_size=num_gpus),
|
|
offload=offload,
|
|
quantization=quantization,
|
|
# The NVFP4 decoder installs its own fused CUDA graphs; compiling the VAE
|
|
# on top trips an inductor cudagraph assertion at 768p with the 36-block student.
|
|
compile=CompileConfig(),
|
|
),
|
|
pipeline=PipelineSelection(
|
|
workload_type="i2v",
|
|
# FastVideo reads the trained fused-block partition and attention
|
|
# settings from the composed directory's fastvideo_inference.json.
|
|
components=ComponentConfig(override_pipeline_cls_name="MiniMaxH3Ref2VAModularPipeline"),
|
|
vae_tiling=profile is not None,
|
|
experimental={
|
|
**experimental, "video_decode_backend": "h3-vae"
|
|
} if profile is not None else experimental,
|
|
),
|
|
)
|
|
|
|
|
|
def main() -> None:
|
|
args = parse_args()
|
|
output_dir = Path(args.output)
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
model_dir, contract = resolve_model(args)
|
|
print(f"Composed model directory: {model_dir}")
|
|
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, reference keep rate {contract['vsa_ref_keep_rate']}")
|
|
references = [MiniMaxH3Reference(source=path, media_type=kind) for kind, path in args.references]
|
|
|
|
profile = DEVICE_PROFILES[args.profile] if args.profile else None
|
|
if profile is not None:
|
|
apply_profile(profile, args.decoder, args.adaln_table)
|
|
print(f"Profile: {profile.label}")
|
|
elif args.lossless_accel:
|
|
apply_lossless_accel_env()
|
|
apply_attention_prefix(args.attention_prefix)
|
|
if args.width is None:
|
|
args.width = profile.width if profile is not None else DEFAULT_WIDTH
|
|
if args.height is None:
|
|
args.height = profile.height if profile is not None else DEFAULT_HEIGHT
|
|
generator = VideoGenerator.from_config(
|
|
build_generator_config(model_dir, args.num_gpus, args.lossless_accel, profile))
|
|
try:
|
|
result = generator.generate(
|
|
GenerationRequest(
|
|
prompt=args.prompt,
|
|
negative_prompt="",
|
|
inputs=InputConfig(references=references),
|
|
sampling=SamplingConfig(
|
|
height=args.height,
|
|
width=args.width,
|
|
num_frames=args.num_frames,
|
|
fps=24,
|
|
# PDD counts fused blocks: one transformer forward each.
|
|
num_inference_steps=contract["num_inference_steps"],
|
|
guidance_scale=1.0,
|
|
batch_cfg=False,
|
|
seed=args.seed,
|
|
),
|
|
output=OutputConfig(
|
|
output_path=str(output_dir / "fasth3_omniref_pdd.mp4"),
|
|
save_video=True,
|
|
return_frames=False,
|
|
),
|
|
))
|
|
print(f"Output written to: {result.video_path}")
|
|
if result.generation_time is not None:
|
|
print(f"Generation time: {result.generation_time:.1f} s")
|
|
if result.peak_memory_mb is not None:
|
|
print(f"Peak GPU memory allocated (rank 0): {result.peak_memory_mb / 1024:.1f} GiB")
|
|
finally:
|
|
generator.shutdown()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|