Spaces:
Running on Zero
Running on Zero
| # 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", | |
| } | |
| 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() | |