# 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()