Optimization Examples¶

"""Fast NVFP4 linear inference for Wan2.1-T2V-1.3B with TAEHV decoding.

This is the FP4-linear fast path from ``fp4_linear_wan2_1_1_3b.py`` with the
heavy Wan VAE swapped out for TAEHV -- a tiny autoencoder that decodes Wan2.1
latents directly (no denormalization) and is dramatically faster / lighter.

How it works: the generator runs with ``output_type="latent"`` so the pipeline
returns raw denoised latents instead of pixels (the Wan VAE is offloaded and
never used). We then decode those latents with TAEHV in this script and save
the frames ourselves. This mirrors the FastVideo-Quantization
``quantization_example_taehv.py`` proof-of-concept, but kept clean: TAEHV is a
pip package (no ``sys.path`` hacks), the latent->uint8 conversion is vectorized,
and there is no dead profiler / sanitization code.

Requirements:
    - Blackwell GPU (B200/B300, sm100a/sm103a) for the FP4 linear path
    - flashinfer (``pip install flashinfer-python``)
    - TAEHV weights ``taew2_1.pth`` (https://github.com/madebyollin/taehv)

Usage:
    python fp4_linear_taehv_wan2_1_1_3b.py                    # FP4 + TAEHV + compile
    python fp4_linear_taehv_wan2_1_1_3b.py --no-taehv         # FP4 + full Wan VAE
    python fp4_linear_taehv_wan2_1_1_3b.py --no-compile       # eager
    python fp4_linear_taehv_wan2_1_1_3b.py --baseline         # dense bf16 reference
    python fp4_linear_taehv_wan2_1_1_3b.py --distilled_model ''  # base Wan2.1 weights
    python fp4_linear_taehv_wan2_1_1_3b.py --warmups 5 --benchmark-runs 20  # timing stats (default)
"""

import argparse
import contextlib
import logging
import os
import statistics
import time

import imageio
import torch

from fastvideo import VideoGenerator
from fastvideo.configs.pipelines.base import PipelineConfig
from fastvideo.layers.quantization.nvfp4_qat_config import NVFP4QATConfig

OUTPUT_PATH = "video_samples"

# Distilled, quantization-aware (QAD) transformer for Wan2.1-1.3B (3 steps,
# guidance 1.0). Loaded on top of the base Wan2.1 pipeline; pass
# ``--distilled_model ''`` to run the base weights instead.
DEFAULT_DISTILLED_MODEL = "FastVideo/FastWan-QAD-1.3B"
DISTILLED_WEIGHTS_FILE = ("generator_inference_transformer/diffusion_pytorch_model.safetensors")

# TAEHV checkpoint for Wan2.1. Clone https://github.com/madebyollin/taehv to get
# ``taew2_1.pth`` (Wan 2.1 / Wan 2.2-14B / Qwen-Image all use this VAE).
DEFAULT_TAEHV_CHECKPOINT = "/root/taehv/taew2_1.pth"

PROMPT = ("A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes "
          "wide with interest. The playful yet serene atmosphere is complemented by soft "
          "natural light filtering through the petals. Mid-shot, warm and cheerful tones.")


class TaehvDecoder:
    """Thin wrapper around the TAEHV tiny autoencoder for Wan2.1 latents.

    TAEHV consumes the *normalized* latents the diffusion model produces (the
    same representation FastVideo carries internally), so no denormalization is
    needed -- unlike the full Wan VAE path.
    """

    def __init__(self, checkpoint_path: str, device: str = "cuda", dtype: torch.dtype = torch.float16) -> None:
        from taehv import TAEHV  # pip-installed; no sys.path manipulation
        self.device = device
        self.dtype = dtype
        print(f"Loading TAEHV from {checkpoint_path} ...")
        self.model = TAEHV(checkpoint_path=checkpoint_path).to(device, dtype).eval()

    @torch.no_grad()
    def decode(self, latents: torch.Tensor):
        """Decode FastVideo latents into uint8 RGB frames.

        Args:
            latents: ``[B, C, T, H, W]`` (NCTHW) normalized latent tensor.

        Returns:
            A ``(T, H, W, 3)`` uint8 numpy array ready for ``imageio.mimsave``.
        """
        # NCTHW -> NTCHW (TAEHV's expected layout), on the TAEHV device/dtype.
        latents = latents.permute(0, 2, 1, 3, 4).to(self.device, self.dtype)
        decoded = self.model.decode_video(latents, parallel=True, show_progress_bar=False)
        # decoded: [B, T, 3, H, W] in [0, 1]. Take batch 0, vectorize to uint8.
        frames = (decoded[0].clamp(0, 1) * 255).to(torch.uint8)
        return frames.permute(0, 2, 3, 1).cpu().numpy()


def resolve_distilled_weights(hf_id: str) -> str:
    """Return a local path to the distilled transformer safetensors."""
    if os.path.exists(hf_id):
        return hf_id
    from huggingface_hub import hf_hub_download
    return hf_hub_download(repo_id=hf_id, filename=DISTILLED_WEIGHTS_FILE)


@contextlib.contextmanager
def silence_request_log():
    """Quiet ``VideoGenerator.generate``'s per-request config printout.

    Each ``generate(...)`` call logs a multi-line debug block (height/width/
    prompt/steps/...) at INFO via ``logger.info`` in
    ``fastvideo.entrypoints.video_generator``. There is no built-in switch,
    so this context manager raises that logger's level to WARNING while the
    warmup calls run, then restores it for the timed run.
    """
    vg_logger = logging.getLogger("fastvideo.entrypoints.video_generator")
    prev_level = vg_logger.level
    vg_logger.setLevel(logging.WARNING)
    try:
        yield
    finally:
        vg_logger.setLevel(prev_level)


def resolve_taehv_checkpoint(path: str) -> str:
    """Validate the TAEHV checkpoint path, with a helpful error if missing."""
    if os.path.exists(path):
        return path
    raise FileNotFoundError(f"TAEHV checkpoint not found at {path!r}. Clone the weights with:\n"
                            "    git clone https://github.com/madebyollin/taehv\n"
                            "and pass --taehv_checkpoint <repo>/taew2_1.pth")


def build_generator(args: argparse.Namespace) -> VideoGenerator:
    model_id = args.model

    # Half precision everywhere; DiT linears are additionally NVFP4-quantized
    # via dit_config.quant_config below.
    pipeline_config = PipelineConfig.from_pretrained(model_id)
    pipeline_config.dit_precision = "bf16"
    pipeline_config.vae_precision = "bf16"
    pipeline_config.text_encoder_precisions = ("bf16", )

    if not args.baseline:
        pipeline_config.dit_config.quant_config = NVFP4QATConfig()

    compile_enabled = not args.no_compile

    extra_kwargs = {}
    if args.distilled_model:
        weights_path = resolve_distilled_weights(args.distilled_model)
        print(f"Using distilled weights: {args.distilled_model} -> {weights_path}")
        extra_kwargs["init_weights_from_safetensors"] = weights_path

    if args.taehv:
        # Skip the in-pipeline VAE decode entirely: the pipeline returns raw
        # latents, the Wan VAE is offloaded to CPU (and not compiled) since we
        # decode with TAEHV in this script instead.
        extra_kwargs["output_type"] = "latent"

    generator = VideoGenerator.from_pretrained(
        model_id,
        pipeline_config=pipeline_config,
        num_gpus=args.num_gpus,
        # Keep everything resident on the GPU -- no offloading, except the
        # unused Wan VAE when TAEHV handles decoding.
        use_fsdp_inference=False,
        dit_cpu_offload=False,
        dit_layerwise_offload=False,
        vae_cpu_offload=args.taehv,
        text_encoder_cpu_offload=False,
        pin_cpu_memory=False,
        enable_torch_compile=compile_enabled,
        enable_torch_compile_text_encoder=compile_enabled,
        enable_torch_compile_vae=compile_enabled and not args.taehv,
        **extra_kwargs,
    )
    return generator


def main() -> None:
    parser = argparse.ArgumentParser(description="FP4 linear Wan2.1-1.3B with TAEHV decoding benchmark")
    parser.add_argument("--model", default="Wan-AI/Wan2.1-T2V-1.3B-Diffusers", help="Model path or HuggingFace ID")
    parser.add_argument("--baseline", action="store_true", help="Run dense bf16 instead of FP4 linear")
    parser.add_argument("--no-compile", action="store_true", help="Disable torch.compile (eager)")
    parser.add_argument("--taehv",
                        action=argparse.BooleanOptionalAction,
                        default=True,
                        help="Decode with TAEHV instead of the full Wan VAE "
                        "(use --no-taehv for the Wan VAE path)")
    parser.add_argument("--taehv_checkpoint",
                        default=DEFAULT_TAEHV_CHECKPOINT,
                        help="Path to the TAEHV taew2_1.pth checkpoint")
    parser.add_argument("--distilled_model",
                        default=DEFAULT_DISTILLED_MODEL,
                        help="HuggingFace ID (or local path) of a distilled "
                        "transformer checkpoint to load on top of --model. "
                        "Pass '' to use the base --model weights instead.")
    parser.add_argument("--num_gpus", type=int, default=1)
    parser.add_argument("--infer_steps", type=int, default=3)
    parser.add_argument("--guidance_scale", type=float, default=1.0)
    parser.add_argument("--warmups", type=int, default=5, help="Warmup runs before timing (default: 5).")
    parser.add_argument("--benchmark-runs",
                        type=int,
                        default=20,
                        help="Timed runs to collect min/max/mean/std over (default: 20).")
    args = parser.parse_args()

    if not torch.cuda.is_available():
        raise SystemExit("CUDA is required for FP4 inference.")

    cap = torch.cuda.get_device_capability()
    print(f"GPU: {torch.cuda.get_device_name()} (capability {cap[0]}.{cap[1]})")
    if not args.baseline and cap[0] < 10:
        print("Warning: NVFP4 requires Blackwell (capability 10.0+); "
              "FP4 kernels may be unavailable on this GPU.")

    mode = "bf16" if args.baseline else "fp4_linear"
    mode += "_taehv" if args.taehv else "_wanvae"
    if not args.no_compile:
        mode += "_compile"
    print(f"Mode: {mode.upper()}")

    # Load TAEHV before the (slow) generator build so a bad checkpoint path
    # fails fast.
    taehv = TaehvDecoder(resolve_taehv_checkpoint(args.taehv_checkpoint)) \
        if args.taehv else None

    generator = build_generator(args)

    os.makedirs(OUTPUT_PATH, exist_ok=True)

    # Warmup: pay the DiT torch.compile cost and warm TAEHV's cuDNN algo
    # selection + allocator growth (~0.2s on the first decode) so the timed
    # runs below measure steady-state latency only.
    with silence_request_log():
        for _ in range(args.warmups):
            warm = generator.generate(
                request={
                    "prompt": PROMPT,
                    "sampling": {
                        "num_inference_steps": 2,
                        "guidance_scale": args.guidance_scale
                    },
                    "output": {
                        "save_video": False,
                        "return_frames": args.taehv
                    },
                })
            if args.taehv:
                taehv.decode(warm.samples)

    output_path = os.path.join(OUTPUT_PATH, f"raccoon_{mode}.mp4")

    # Benchmark: time each run end-to-end. ``denoise`` is the generator's own
    # generation_time; for TAEHV we add the in-script decode. The mp4 is not
    # written inside the loop (only the final run's frames are saved below) so
    # disk I/O never pollutes the timings.
    denoise_times: list[float] = []
    decode_times: list[float] = []
    totals: list[float] = []
    frames = None
    with silence_request_log():
        for i in range(args.benchmark_runs):
            result = generator.generate(
                request={
                    "prompt": PROMPT,
                    "sampling": {
                        "num_inference_steps": args.infer_steps,
                        "guidance_scale": args.guidance_scale,
                    },
                    "output": {
                        "save_video": False,
                        "return_frames": args.taehv,
                        "output_path": output_path,
                    },
                })
            denoise_elapsed = result.generation_time
            denoise_times.append(denoise_elapsed)

            if args.taehv:
                torch.cuda.synchronize()
                decode_start = time.perf_counter()
                frames = taehv.decode(result.samples)
                torch.cuda.synchronize()
                decode_elapsed = time.perf_counter() - decode_start
                decode_times.append(decode_elapsed)
                total = denoise_elapsed + decode_elapsed
            else:
                total = denoise_elapsed
            totals.append(total)

            line = f"  run {i + 1:02d}/{args.benchmark_runs}: {total:.3f}s"
            if args.taehv:
                line += (f" (denoise {denoise_elapsed:.3f}s + "
                         f"decode {decode_elapsed:.3f}s)")
            print(line)

    if args.taehv and frames is not None:
        imageio.mimsave(output_path, frames, fps=16, format="mp4")
        print(f"Saved video to {output_path}")

    # Report min / max / mean / std over the timed runs.
    def _stat_row(name: str, xs: list[float]) -> str:
        std = statistics.stdev(xs) if len(xs) > 1 else 0.0
        return (f"  {name:<11}{min(xs):>8.3f}{max(xs):>9.3f}"
                f"{statistics.mean(xs):>9.3f}{std:>9.3f}")

    if totals:
        print(f"\n[{mode.upper()}] {args.benchmark_runs} runs, {args.warmups} warmup, "
              f"{args.infer_steps} steps:")
        print(f"  {'metric':<11}{'min':>8}{'max':>9}{'mean':>9}{'std':>9}   (s)")
        print(_stat_row("total", totals))
        if args.taehv:
            print(_stat_row("denoise", denoise_times))
            print(_stat_row("decode", decode_times))

    generator.shutdown()


if __name__ == "__main__":
    main()

import os
import time

from fastvideo import VideoGenerator


def main():
    # set the attention backend
    os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "FLASH_ATTN"

    start_time = time.perf_counter()
    gen = VideoGenerator.from_pretrained(
        model_path="Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
        num_gpus=1,
        dit_cpu_offload=False,
        vae_cpu_offload=False,
        text_encoder_cpu_offload=True,
        pin_cpu_memory=True,  # set to false if low CPU RAM or hit obscure "CUDA error: Invalid argument"
    )
    load_time = time.perf_counter() - start_time
    print(f"Model loading time: {load_time:.2f} seconds")

    gen_start_time = time.perf_counter()

    gen.generate_video(
        prompt=
        "Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of a bustling street food market. The scene captures a mix of humor and authenticity. Mid-shot framing, vibrant lighting.",
        seed=1024,
        output_path="example_outputs/")

    generation_time = time.perf_counter() - gen_start_time
    print(f"Video generation time: {generation_time:.2f} seconds")

    total_time = time.perf_counter() - start_time
    print(f"Total execution time: {total_time:.2f} seconds")


if __name__ == "__main__":
    main()

"""FP4 Flash Attention 4 inference example on Blackwell GPUs.

Quantizes Q and K to NVFP4 E2M1 with per-block E4M3 scale factors,
achieving up to 1.39x attention kernel speedup over BF16 FA4.

Requirements:
    - Blackwell GPU (B200/B300, sm100a/sm103a)
    - flash-attention-fp4, cutlass-dsl, flashinfer
    - See docs/inference/optimizations.md for installation

Usage:
    python fp4_attn_wan2_1_1_3b.py --nvfp4_fa4
    python fp4_attn_wan2_1_1_3b.py  # BF16 baseline
"""

import argparse
import os
import time

from fastvideo import VideoGenerator

OUTPUT_PATH = "video_samples"


def main():
    parser = argparse.ArgumentParser(description="FP4 FA4 video generation benchmark")
    parser.add_argument("--nvfp4_fa4", action="store_true", help="Enable NVFP4 FP4 quantized QK flash attention")
    parser.add_argument("--model", default="Wan-AI/Wan2.1-T2V-1.3B-Diffusers", help="Model path or HuggingFace ID")
    parser.add_argument("--compile", action="store_true", help="Enable torch.compile for DIT")
    parser.add_argument("--num_gpus", type=int, default=1)
    parser.add_argument("--infer_steps", type=int, default=50)
    args = parser.parse_args()

    mode = "nvfp4" if args.nvfp4_fa4 else "bf16"
    if args.compile:
        mode += "_compile"
    print(f"Mode: {mode.upper()}")

    generator = VideoGenerator.from_pretrained(
        args.model,
        num_gpus=args.num_gpus,
        nvfp4_fa4=args.nvfp4_fa4,
        use_fsdp_inference=not args.nvfp4_fa4,
        dit_cpu_offload=False,
        dit_layerwise_offload=False,
        vae_cpu_offload=True,
        text_encoder_cpu_offload=True,
        enable_torch_compile=args.compile,
    )

    prompt = ("A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes "
              "wide with interest. The playful yet serene atmosphere is complemented by soft "
              "natural light filtering through the petals. Mid-shot, warm and cheerful tones.")

    n_warmup = 2 if args.compile else 1
    for i in range(n_warmup):
        generator.generate(request={
            "prompt": prompt,
            "sampling": {
                "num_inference_steps": 2
            },
            "output": {
                "save_video": False
            }
        })

    os.makedirs(OUTPUT_PATH, exist_ok=True)
    start = time.time()
    generator.generate(
        request={
            "prompt": prompt,
            "sampling": {
                "num_inference_steps": args.infer_steps
            },
            "output": {
                "save_video": True,
                "output_path": os.path.join(OUTPUT_PATH, f"raccoon_{mode}.mp4")
            },
        })
    elapsed = time.time() - start
    print(f"[{mode.upper()}] {args.infer_steps} steps in {elapsed:.2f}s "
          f"({args.infer_steps / elapsed:.2f} it/s)")

    generator.shutdown()


if __name__ == "__main__":
    main()

"""FP8 weight quantization inference example.

Runs Wan2.1-T2V-1.3B with FP8 e4m3 quantized DiT linear layers (attention
projections and FFN). Weights are quantized in-place after loading; activations
are quantized dynamically at runtime. Reduces GPU memory relative to BF16 and
can improve throughput on sm89+ GPUs.

Requirements:
    - GPU: sm89+ (H100, L40S, RTX 4090, Ada Lovelace, or newer)
      Falls back to a bf16 dequant path on older GPUs.
    - TAEHV (optional): Follow install instructions at https://github.com/madebyollin/taehv

Usage:
    python fp8_wan2_1_1_3b.py              # FP8 per-tensor (default)
    python fp8_wan2_1_1_3b.py --bf16       # BF16 baseline
    python fp8_wan2_1_1_3b.py --granularity channel  # per-channel (higher accuracy but slower)
    python fp8_wan2_1_1_3b.py --taehv-checkpoint /path/to/taew2_1.pth
"""

import argparse
import os
import sys
import time

import torch

OUTPUT_PATH = "video_samples"


def load_taehv(checkpoint_path, device="cuda", dtype=torch.float16):
    repo_dir = os.path.dirname(checkpoint_path)
    if repo_dir not in sys.path:
        sys.path.insert(0, repo_dir)
    from taehv import TAEHV
    print(f"Loading TAEHV from {checkpoint_path}...")
    model = TAEHV(checkpoint_path=checkpoint_path).to(device, dtype)
    print("TAEHV loaded.")
    return model


@torch.no_grad()  # type: ignore[misc]
def decode_with_taehv(taehv_model, latents):
    latents = latents.permute(0, 2, 1, 3, 4)
    latents = latents.to(device=next(taehv_model.parameters()).device, dtype=next(taehv_model.parameters()).dtype)
    decoded = taehv_model.decode_video(latents, parallel=False, show_progress_bar=False)
    frames = []
    for frame in decoded[0]:
        frame_np = (frame.clamp(0, 1) * 255).byte().cpu().permute(1, 2, 0).numpy()
        frames.append(frame_np)
    return frames


def main():
    parser = argparse.ArgumentParser(description="FP8 video generation benchmark")
    parser.add_argument("--bf16", action="store_true", help="BF16 baseline (no FP8 quantization)")
    parser.add_argument("--granularity",
                        choices=["tensor", "channel"],
                        default="tensor",
                        help="FP8 weight scale granularity: tensor (faster) or channel (more accurate)")
    parser.add_argument("--taehv-checkpoint",
                        default=None,
                        metavar="PATH",
                        help="Path to taew2_1.pth; enables TAEHV tiny autoencoder decoding")
    parser.add_argument("--model", default="FastVideo/FastWan-QAD-FP8-1.3B", help="Model path or HuggingFace ID")
    parser.add_argument("--no-compile", action="store_true", help="Disable torch.compile for the DiT")
    parser.add_argument("--num_gpus", type=int, default=1)
    parser.add_argument("--infer_steps", type=int, default=3)
    args = parser.parse_args()

    os.environ.setdefault("FASTVIDEO_ATTENTION_BACKEND", "SAGE_ATTN")

    from fastvideo import VideoGenerator
    from fastvideo.layers.quantization import get_quantization_config

    mode = "bf16" if args.bf16 else f"fp8_{args.granularity}"
    if not args.no_compile:
        mode += "_compile"
    use_taehv = args.taehv_checkpoint is not None
    print(f"Mode: {mode.upper()}" + ("  decoder=TAEHV" if use_taehv else "  decoder=VAE"))

    taehv_model = load_taehv(args.taehv_checkpoint) if use_taehv else None

    # transformer_quant needs a QuantizationConfig *instance* — the bare string
    # is not resolved on the from_pretrained kwarg path.
    extra = {} if args.bf16 else {"transformer_quant": get_quantization_config("FP8")(granularity=args.granularity)}
    generator = VideoGenerator.from_pretrained(
        args.model,
        num_gpus=args.num_gpus,
        use_fsdp_inference=False,
        dit_cpu_offload=False,
        dit_layerwise_offload=False,
        vae_cpu_offload=use_taehv,
        text_encoder_cpu_offload=False,
        pin_cpu_memory=False,
        enable_torch_compile=not args.no_compile,
        enable_torch_compile_vae=not args.no_compile and not use_taehv,
        output_type="latent" if use_taehv else "pil",
        **extra,
    )

    prompt = ("A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes "
              "wide with interest. The playful yet serene atmosphere is complemented by soft "
              "natural light filtering through the petals. Mid-shot, warm and cheerful tones.")

    n_warmup = 1 if not args.no_compile else 0
    for _ in range(n_warmup):
        generator.generate(
            request={
                "prompt": prompt,
                "sampling": {
                    "num_inference_steps": 3,
                    "guidance_scale": 1.0
                },
                "output": {
                    "save_video": False
                }
            })

    os.makedirs(OUTPUT_PATH, exist_ok=True)
    start = time.time()
    if use_taehv:
        result = generator.generate(
            request={
                "prompt": prompt,
                "sampling": {
                    "num_inference_steps": args.infer_steps,
                    "guidance_scale": 1.0
                },
                "output": {
                    "save_video": False
                },
            })
        import imageio
        frames = decode_with_taehv(taehv_model, result.samples)
        video_path = os.path.join(OUTPUT_PATH, f"raccoon_{mode}.mp4")
        imageio.mimsave(video_path, frames, fps=16, format="mp4")
        print(f"Saved TAEHV-decoded video to: {video_path}")
    else:
        generator.generate(
            request={
                "prompt": prompt,
                "sampling": {
                    "num_inference_steps": args.infer_steps,
                    "guidance_scale": 1.0
                },
                "output": {
                    "save_video": True,
                    "output_path": os.path.join(OUTPUT_PATH, f"raccoon_{mode}.mp4")
                },
            })
    elapsed = time.time() - start
    print(f"[{mode.upper()}] {args.infer_steps} steps in {elapsed:.2f}s "
          f"({args.infer_steps / elapsed:.2f} it/s)")

    generator.shutdown()


if __name__ == "__main__":
    main()

"""NVFP4 + Attn-QAT (modified SageAttention3) inference on Blackwell.

Runs Wan2.1-T2V-1.3B fully in 4-bit: NVFP4 linear layers (activations
quantized on the fly) together with the modified SageAttention3 FP4 attention
backend (``ATTN_QAT_INFER``). This is the inference half of the
Quantization-Aware Distillation (QAD) recipe.

Requirements:
    - RTX 5090 / consumer Blackwell (sm_120a). The attn_qat_infer kernel hard
      gates on sm_120; on other GPUs it falls back to Flash Attention.
    - The attn_qat_infer kernel built into fastvideo-kernel (see #1455) and
      flashinfer for the NVFP4 linear matmuls.

Usage:
    python nvfp4_qat_wan2_1_1_3b.py                 # NVFP4 linear + Attn-QAT attn
    python nvfp4_qat_wan2_1_1_3b.py --bf16          # BF16 baseline
"""

import argparse
import os
import time

OUTPUT_PATH = "video_samples"


def main():
    parser = argparse.ArgumentParser(description="NVFP4 + Attn-QAT video generation")
    parser.add_argument("--bf16", action="store_true", help="BF16 baseline (no NVFP4 linear, default attention)")
    parser.add_argument("--model", default="Wan-AI/Wan2.1-T2V-1.3B-Diffusers", help="Model path or HuggingFace ID")
    parser.add_argument("--quant-method",
                        default="nvfp4_qat",
                        choices=["nvfp4_qat", "NVFP4"],
                        help="Linear quantization config. Wan-2.1 uses nvfp4_qat (matches its "
                        "to_q/k/v/out + ffn layers); NVFP4 is LTX2-specific and will NOT "
                        "quantize Wan.")
    parser.add_argument("--compile", action="store_true", help="Enable torch.compile for the DiT")
    parser.add_argument("--num_gpus", type=int, default=1)
    parser.add_argument("--infer_steps", type=int, default=50)
    args = parser.parse_args()

    # The attention backend is selected via env var before the engine starts.
    # ATTN_QAT_INFER -> AttnQatInferBackend (modified SageAttention3 FP4).
    if not args.bf16:
        os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "ATTN_QAT_INFER"

    # Import after the env var so the platform picks up the selection.
    from fastvideo import VideoGenerator
    from fastvideo.layers.quantization import get_quantization_config

    mode = "bf16" if args.bf16 else args.quant_method
    if args.compile:
        mode += "_compile"
    print(f"Mode: {mode.upper()}")

    # transformer_quant needs a QuantizationConfig *instance* — the bare string
    # is not resolved on the from_pretrained kwarg path.
    extra = {} if args.bf16 else {"transformer_quant": get_quantization_config(args.quant_method)()}
    generator = VideoGenerator.from_pretrained(
        args.model,
        num_gpus=args.num_gpus,
        use_fsdp_inference=args.bf16,
        dit_cpu_offload=False,
        vae_cpu_offload=True,
        text_encoder_cpu_offload=True,
        enable_torch_compile=args.compile,
        **extra,
    )

    prompt = ("A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes "
              "wide with interest. The playful yet serene atmosphere is complemented by soft "
              "natural light filtering through the petals. Mid-shot, warm and cheerful tones.")

    n_warmup = 2 if args.compile else 1
    for _ in range(n_warmup):
        generator.generate(request={
            "prompt": prompt,
            "sampling": {
                "num_inference_steps": 2
            },
            "output": {
                "save_video": False
            }
        })

    os.makedirs(OUTPUT_PATH, exist_ok=True)
    start = time.time()
    generator.generate(
        request={
            "prompt": prompt,
            "sampling": {
                "num_inference_steps": args.infer_steps
            },
            "output": {
                "save_video": True,
                "output_path": os.path.join(OUTPUT_PATH, f"raccoon_{mode}.mp4")
            },
        })
    elapsed = time.time() - start
    print(f"[{mode.upper()}] {args.infer_steps} steps in {elapsed:.2f}s "
          f"({args.infer_steps / elapsed:.2f} it/s)")

    generator.shutdown()


if __name__ == "__main__":
    main()

"""QAD FP4 quality A/B/C/D harness — Wan2.1-T2V-1.3B on sm_121 (DGX Spark GB10).

Tests whether the quantization-aware-distilled checkpoint
``FastVideo/FastWan-QAD-1.3B`` recovers FP4 quality on sm_121, using the
sm_121-enabled ``ATTN_QAT_INFER`` attention kernel. Running FP4 attention on
*stock* Wan weights gives output below bf16 — expected, because stock weights
were never trained to tolerate FP4 attention. This harness runs the checkpoint
that *was* (fake-quant FP4 attention + NVFP4 linear trained into the model).

Attention (bf16 vs FP4) and linear (bf16 vs FP4) are fully decoupled at
inference, so we can isolate each axis:

    FASTVIDEO_ATTENTION_BACKEND unset          -> bf16 attention (SDPA)
    FASTVIDEO_ATTENTION_BACKEND=ATTN_QAT_INFER -> FP4 attention
    QAD_LINEAR=0                               -> bf16 linear
    QAD_LINEAR=1                               -> NVFP4 FP4 linear (flashinfer)

    arm  attention   linear   selects with
    A    bf16        bf16     QAD_LINEAR=0  (no ATTN env)            reference
    B    FP4         bf16     QAD_LINEAR=0  ATTN_QAT_INFER           isolate attn
    C    bf16        FP4      QAD_LINEAR=1  (no ATTN env)            isolate linear
    D    FP4         FP4      QAD_LINEAR=1  ATTN_QAT_INFER           full 4-bit

All four arms run end-to-end on the GB10 (the full 4-bit path — FP4 linear +
FP4 attention — works on sm_121). This script still runs exactly ONE arm per
invocation and dumps a C stack on any hard crash, so a single misbehaving arm
can never take the others down with it; the runbook loops it four times with
different env. Quality is the eye/ear on the saved mp4 + a matching-frame still;
timing is the mean generation_time (full pipeline: text-encode +
denoise + VAE decode) over the measured runs.

On the GB10, expect FP4 attention ~6% faster end-to-end generation vs bf16 and quality-neutral
by eye on the QAD checkpoint (both share the 3-step distill's quality ceiling).
FP4 *linear* is roughly break-even at 1.3B/480p (the per-call quantize overhead
eats the small-GEMM saving in eager mode); its win shows at higher resolution
with torch.compile.

Knobs (env): QAD_MODEL, QAD_DISTILLED, QAD_STEPS (3), QAD_GUIDANCE (1.0),
QAD_HEIGHT (480), QAD_WIDTH (832), QAD_FRAMES (77), QAD_SEED (42),
QAD_WARMUP (1), QAD_RUNS (3), QAD_STILL (20), QAD_OUT (qad_fp4_samples).
"""
from __future__ import annotations

import faulthandler
import glob
import os
import time

import torch

faulthandler.enable()  # dump a C stack if any arm hard-crashes.


def _env(name: str, default: str) -> str:
    return os.environ.get(name, default)


def _env_int(name: str, default: int) -> int:
    return int(os.environ.get(name, str(default)))


def _env_float(name: str, default: float) -> float:
    return float(os.environ.get(name, str(default)))


# Distilled QAD transformer, loaded on top of the base Wan2.1-1.3B pipeline.
DEFAULT_DISTILLED = "FastVideo/FastWan-QAD-1.3B"
# The repo is a full diffusers pipeline; we overlay only its transformer onto the
# base Wan pipeline (vae/text_encoder are Wan-identical).
DISTILLED_WEIGHTS_FILE = "transformer/diffusion_pytorch_model.safetensors"

PROMPT = ("A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes "
          "wide with interest. The playful yet serene atmosphere is complemented by soft "
          "natural light filtering through the petals. Mid-shot, warm and cheerful tones.")


def resolve_distilled_weights(hf_id: str) -> str:
    """Return a local path to the distilled transformer safetensors."""
    if not hf_id:
        return ""
    if os.path.exists(hf_id):
        return hf_id
    from huggingface_hub import hf_hub_download
    return hf_hub_download(repo_id=hf_id, filename=DISTILLED_WEIGHTS_FILE)


def build_generator(fp4_linear: bool):
    from fastvideo import VideoGenerator
    from fastvideo.configs.pipelines.base import PipelineConfig

    model_id = _env("QAD_MODEL", "Wan-AI/Wan2.1-T2V-1.3B-Diffusers")

    pipeline_config = PipelineConfig.from_pretrained(model_id)
    pipeline_config.dit_precision = "bf16"
    pipeline_config.vae_precision = "bf16"
    pipeline_config.text_encoder_precisions = ("bf16", )

    if fp4_linear:
        # Wan-style config: matches to_q/k/v/out + ffn (the plain NVFP4 config is
        # LTX2-specific and would quantize nothing on Wan).
        from fastvideo.layers.quantization.nvfp4_qat_config import NVFP4QATConfig
        pipeline_config.dit_config.quant_config = NVFP4QATConfig()

    extra_kwargs = {}
    distilled = resolve_distilled_weights(_env("QAD_DISTILLED", DEFAULT_DISTILLED))
    if distilled:
        print(f"[qad] distilled weights: {distilled}")
        extra_kwargs["init_weights_from_safetensors"] = distilled

    # Keep everything resident (the 1.3B QAD model + FP4 fits the GB10's unified
    # memory); real Wan VAE decode for a faithful quality read (no TAEHV).
    return VideoGenerator.from_pretrained(
        model_id,
        pipeline_config=pipeline_config,
        num_gpus=1,
        use_fsdp_inference=False,
        dit_cpu_offload=False,
        dit_layerwise_offload=False,
        vae_cpu_offload=False,
        text_encoder_cpu_offload=False,
        pin_cpu_memory=False,
        enable_torch_compile=False,  # eager: isolate the FP4 effect, no compile noise
        **extra_kwargs,
    )


def main() -> None:
    if not torch.cuda.is_available():
        raise SystemExit("CUDA is required.")

    fp4_linear = _env_int("QAD_LINEAR", 0) == 1
    attn = _env("FASTVIDEO_ATTENTION_BACKEND", "default")
    steps = _env_int("QAD_STEPS", 3)
    guidance = _env_float("QAD_GUIDANCE", 1.0)
    height = _env_int("QAD_HEIGHT", 480)
    width = _env_int("QAD_WIDTH", 832)
    frames = _env_int("QAD_FRAMES", 77)
    seed = _env_int("QAD_SEED", 42)
    warmup = _env_int("QAD_WARMUP", 1)
    runs = _env_int("QAD_RUNS", 3)
    still_idx = _env_int("QAD_STILL", 20)
    out_dir = _env("QAD_OUT", "qad_fp4_samples")
    prompt = _env("QAD_PROMPT", PROMPT)

    tag = f"lin-{'fp4' if fp4_linear else 'bf16'}_attn-{attn.lower()}"
    cap = torch.cuda.get_device_capability()
    print(f"[qad] GPU {torch.cuda.get_device_name()} (cc {cap[0]}.{cap[1]})")
    print(f"[qad] ARM {tag}: linear={'FP4' if fp4_linear else 'bf16'}, "
          f"attention={attn}, {steps} steps, guidance {guidance}, "
          f"{height}x{width}x{frames}, seed {seed}")
    print(f"[qad] prompt: {prompt[:80]}{'...' if len(prompt) > 80 else ''}")

    arm_dir = os.path.join(out_dir, tag)
    os.makedirs(arm_dir, exist_ok=True)
    generator = build_generator(fp4_linear)

    def _generate():
        # seed + frame dims live under `sampling` (SamplingConfig); `output`
        # only takes output_path/save_video/return_frames (OutputConfig).
        return generator.generate(
            request={
                "prompt": prompt,
                "sampling": {
                    "seed": seed,
                    "num_inference_steps": steps,
                    "guidance_scale": guidance,
                    "height": height,
                    "width": width,
                    "num_frames": frames,
                },
                "output": {
                    "save_video": True,
                    "output_path": arm_dir,
                    "return_frames": True
                },
            })

    for _ in range(warmup):
        _generate()

    denoise_times: list[float] = []
    last = None
    for i in range(runs):
        torch.cuda.synchronize()
        t0 = time.perf_counter()
        last = _generate()
        torch.cuda.synchronize()
        wall = time.perf_counter() - t0
        # generate_video returns a plain dict; attribute access would always
        # fall back to wall time.
        gen_t = last.get("generation_time") if isinstance(last, dict) else None
        denoise_times.append(gen_t if gen_t is not None else wall)
        print(f"[qad] {tag} run {i + 1}/{runs}: {wall:.2f}s wall "
              f"(gen {denoise_times[-1]:.2f}s)")

    # The pipeline wrote the mp4 (full known-good encode) into arm_dir; report
    # it and pull a matching-frame still from the [b,c,t,h,w] samples tensor
    # using the same recipe the pipeline's frame builder uses.
    mp4s = sorted(glob.glob(os.path.join(arm_dir, "*.mp4")), key=os.path.getmtime)
    if mp4s:
        print(f"[qad] video: {mp4s[-1]}")
    samples = getattr(last, "samples", None) if last is not None else None
    if samples is not None and getattr(samples, "ndim", 0) == 5:
        import imageio
        f = min(still_idx, samples.shape[2] - 1)  # samples: [b, c, t, h, w]
        still = (samples[0, :, f].permute(1, 2, 0).clamp(0, 1) * 255)
        still = still.to(torch.uint8).cpu().numpy()
        png = os.path.join(arm_dir, f"raccoon_{tag}_f{f}.png")
        imageio.imwrite(png, still)
        print(f"[qad] still: {png}")
    else:
        print("[qad] note: no 5-D samples tensor; grab a frame from the mp4 above")

    mean = sum(denoise_times) / len(denoise_times)
    print(f"\n[qad][{tag}] generation mean {mean:.2f}s over {runs} runs "
          f"({warmup} warmup, {steps} steps)")
    generator.shutdown()


if __name__ == "__main__":
    main()

"""DGX Spark (GB10) reproduction benchmark for the performance guide.

Reproduces the two headline claims in
``docs/getting_started/installation/spark_performance.md`` on your own GB10:

  1. A distilled few-step model is usable (~30 s/video) and **decode-bound** on
     the Spark's unified LPDDR5X memory.
  2. bf16 VAE decode is essentially lossless (MS-SSIM ~1.0 vs fp32) and modestly
     faster — which is why FastVideo already defaults Wan's *decode* to bf16.

Two parts, both in-process so they control for the Spark's run-to-run variance:

  A. **Generation timing** — loads a distilled model once, generates ``--runs``
     videos back-to-back (after ``--warmup``), reports the *median* generation
     time. Set ``FASTVIDEO_STAGE_LOGGING=1`` to also see the per-stage
     (denoise / VAE decode / text-encode) split that shows the decode bottleneck.

  B. **Decode precision A/B** — decodes ONE fixed latent fp32-vs-bf16 in the same
     process and reports MS-SSIM + speedup. This isolates the decode delta with
     no denoise non-determinism and no video-codec noise.

**Why median, not a single run:** on the GB10 a 3-step generation is dominated by
one-time per-process startup (Triton autotune, allocator warmup) that never
amortizes over so few steps, so single-run totals wobble ~±30%. Always compare
few-step levers back-to-back / as medians, never as two separate single runs.

Run it safely on a shared box (see the best-practices note in the perf guide):

    FASTVIDEO_STAGE_LOGGING=1 nice -n 19 nohup \
        python examples/inference/optimizations/spark_benchmark.py > spark_bench.log 2>&1 &
    tail -f spark_bench.log

Knobs (flags or env): --model, --runs (3), --warmup (1), --steps (3),
--frames (81), --height (448), --width (832), --seed (42), --out, --skip-gen,
--skip-decode.
"""
from __future__ import annotations

import argparse
import os
import statistics
import time

import torch


def _p(msg: str) -> None:
    print(f"[spark-bench] {msg}", flush=True)


def bench_generation(args) -> None:
    """Part A: median few-step generation time on a distilled model."""
    from fastvideo import VideoGenerator
    from fastvideo.api.sampling_param import SamplingParam

    # VSA auto-routes to the Triton kernel on sm_121; do NOT force TORCH_SDPA on a
    # VSA checkpoint (the SDPA path builds a model without the gate weights the
    # checkpoint carries and fails to load).
    os.environ.setdefault("FASTVIDEO_ATTENTION_BACKEND", "VIDEO_SPARSE_ATTN")

    _p(f"loading {args.model} ...")
    load_t0 = time.perf_counter()
    generator = VideoGenerator.from_pretrained(
        args.model,
        num_gpus=1,
        use_fsdp_inference=False,
        # Leave offload at these defaults: "CPU" offload is the same unified RAM
        # on the GB10, so the win is tiling + sane resolution, not offloading.
        text_encoder_cpu_offload=True,
        pin_cpu_memory=True,
        dit_cpu_offload=False,
        vae_cpu_offload=False,
        VSA_sparsity=0.8,
    )
    _p(f"model loaded in {time.perf_counter() - load_t0:.1f}s")

    sampling_param = SamplingParam.from_pretrained(args.model)
    sampling_param.num_frames = args.frames
    sampling_param.height = args.height
    sampling_param.width = args.width
    sampling_param.num_inference_steps = args.steps
    sampling_param.seed = args.seed

    prompt = ("A curious raccoon peers through a vibrant field of yellow sunflowers, "
              "its eyes wide with interest. Soft natural light, warm cheerful tones, "
              "mid-shot, cinematic.")

    def _gen():
        torch.cuda.synchronize()
        t0 = time.perf_counter()
        video = generator.generate_video(prompt, output_path=args.out, save_video=True, sampling_param=sampling_param)
        torch.cuda.synchronize()
        # generate_video returns a plain dict (legacy result), not an object —
        # attribute access would silently fall back to wall time / None.
        dt = video.get("generation_time") if isinstance(video, dict) else None
        if dt is None:
            dt = time.perf_counter() - t0
        # Peak memory is measured *inside the worker process* that runs the
        # pipeline and surfaced on the result; reading torch's allocator in this
        # (main) process would report ~0 because the allocations aren't here.
        peak = video.get("peak_memory_mb") if isinstance(video, dict) else None
        return dt, peak

    for _ in range(args.warmup):
        _gen()

    times, peaks = [], []
    for i in range(args.runs):
        dt, peak = _gen()
        times.append(dt)
        if peak:
            peaks.append(peak)
        _p(f"gen run {i + 1}/{args.runs}: {dt:.2f}s")

    med = statistics.median(times)
    _p(f"median generation time over {args.runs} runs "
       f"({args.warmup} warmup, {args.steps} steps): {med:.2f}s")
    if peaks:
        _p(f"peak GPU memory (worker, reported by pipeline): {max(peaks):.0f} MB "
           f"= {max(peaks) / 1024:.1f} GB")
    else:
        _p("peak GPU memory: not reported by this pipeline build")
    _p("set FASTVIDEO_STAGE_LOGGING=1 to see the denoise / decode / text split "
       "(few-step generation is VAE-decode-bound on the GB10).")
    generator.shutdown()


def bench_decode_precision(args) -> None:
    """Part B: fp32-vs-bf16 VAE decode of one fixed latent (SSIM + speedup)."""
    try:
        from diffusers import AutoencoderKLWan
        from torchmetrics.functional import (multiscale_structural_similarity_index_measure as msssim)
    except ImportError as e:  # torchmetrics is not a hard FastVideo dep
        _p(f"skipping decode A/B (missing dependency: {e}); "
           "`uv pip install torchmetrics` to enable it.")
        return

    dev = "cuda"
    vae = AutoencoderKLWan.from_pretrained(args.model, subfolder="vae", torch_dtype=torch.float32).to(dev).eval()

    # Wan latent geometry for height x width x frames, patch (4,8,8):
    #   T_lat = (frames - 1) // 4 + 1,  H_lat = height // 8,  W_lat = width // 8
    t_lat = (args.frames - 1) // 4 + 1
    z = torch.randn(1, 16, t_lat, args.height // 8, args.width // 8, device=dev, dtype=torch.float32)

    def _decode(autocast: bool):
        torch.cuda.synchronize()
        t0 = time.perf_counter()
        with torch.no_grad(), torch.autocast("cuda", dtype=torch.bfloat16, enabled=autocast):
            out = vae.decode(z, return_dict=False)[0]
        torch.cuda.synchronize()
        return out.float(), time.perf_counter() - t0

    def _frames(o):  # (1,3,T,H,W) [-1,1] -> (T,3,H,W) [0,1]
        return ((o.clamp(-1, 1) + 1) / 2)[0].permute(1, 0, 2, 3).contiguous()

    _decode(False)  # warm both paths (excluded from timing)
    _decode(True)

    o32, t32 = _decode(autocast=False)  # fp32
    o16, t16 = _decode(autocast=True)  # bf16 (== vae_decode_precision="bf16")
    ssim = msssim(_frames(o16), _frames(o32), data_range=1.0).item()

    _p(f"fp32 decode : {t32 * 1000:8.1f} ms")
    _p(f"bf16 decode : {t16 * 1000:8.1f} ms   ({t32 / t16:.2f}x faster)")
    _p(f"MS-SSIM(bf16, fp32) on identical latent: {ssim:.4f}  "
       "(>= ~0.99 -> lossless; this is why Wan decode defaults to bf16)")


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--model", default="FastVideo/FastWan2.1-T2V-1.3B-Diffusers")
    ap.add_argument("--runs", type=int, default=3)
    ap.add_argument("--warmup", type=int, default=1)
    ap.add_argument("--steps", type=int, default=3)
    ap.add_argument("--frames", type=int, default=81)
    ap.add_argument("--height", type=int, default=448)
    ap.add_argument("--width", type=int, default=832)
    ap.add_argument("--seed", type=int, default=42)
    ap.add_argument("--out", default="spark_bench_samples")
    ap.add_argument("--skip-gen", action="store_true", help="skip Part A (generation timing)")
    ap.add_argument("--skip-decode", action="store_true", help="skip Part B (decode precision A/B)")
    args = ap.parse_args()

    if not torch.cuda.is_available():
        raise SystemExit("CUDA is required (run this on the GB10).")
    cap = torch.cuda.get_device_capability()
    _p(f"{torch.cuda.get_device_name()} (cc {cap[0]}.{cap[1]}), "
       f"torch {torch.__version__}")

    if not args.skip_gen:
        bench_generation(args)
    if not args.skip_decode:
        bench_decode_precision(args)


if __name__ == "__main__":
    main()

from fastvideo import VideoGenerator
import argparse

OUTPUT_PATH = "video_samples_wan2_2_5B_ti2v"


def main(text_encoder_path: str):
    # FastVideo will automatically use the optimal default arguments for the
    # model.
    # If a local path is provided, FastVideo will make a best effort
    # attempt to identify the optimal arguments.
    model_name = "Wan-AI/Wan2.2-TI2V-5B-Diffusers"
    generator = VideoGenerator.from_pretrained(
        model_name,
        # FastVideo will automatically handle distributed setup
        num_gpus=1,
        use_fsdp_inference=True,
        dit_cpu_offload=True,
        vae_cpu_offload=False,
        text_encoder_cpu_offload=False,
        # AbsMaxFP8 is the quantization method used by ComfyUI;
        # check fastvideo/layers/quantization/* for more quantization methods
        override_text_encoder_quant="AbsMaxFP8",
        # for Wan 2.2, this is the path to "umt5_xxl_fp8_e4m3fn_scaled.safetensors"
        override_text_encoder_safetensors=text_encoder_path,
        pin_cpu_memory=True,  # set to false if low CPU RAM or hit obscure "CUDA error: Invalid argument"
    )

    # I2V is triggered just by passing in an image_path argument
    prompt = "Summer beach vacation style, a white cat wearing sunglasses sits on a surfboard. The fluffy-furred feline gazes directly at the camera with a relaxed expression. Blurred beach scenery forms the background featuring crystal-clear waters, distant green hills, and a blue sky dotted with white clouds. The cat assumes a naturally relaxed posture, as if savoring the sea breeze and warm sunlight. A close-up shot highlights the feline's intricate details and the refreshing atmosphere of the seaside."
    image_path = "https://huggingface.co/datasets/YiYiXu/testing-images/resolve/main/wan_i2v_input.JPG"
    video = generator.generate_video(prompt, output_path=OUTPUT_PATH, save_video=True, image_path=image_path)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--text_encoder_path",
        type=str,
        required=True,
        help="Path to the quantized text encoder safetensors file.",
    )
    args = parser.parse_args()
    main(args.text_encoder_path)

"""torch.compile A/B example for FastVideo.

`enable_torch_compile=True` compiles the DiT submodules that declare
`_compile_conditions` for a substantial end-to-end speedup (e.g.
Wan2.1-T2V-1.3B on A100: ~-24% e2e). It is off by default.

The first compiled generation pays a one-time graph-build cost; it
amortizes over later generations with the same input shapes. This
script does one un-measured warmup then a measured run so the reported
number is steady-state, not graph-build — measuring the warmup is the
most common way to wrongly conclude compile is slower.

Usage:
    # baseline (eager)
    python torch_compile_example.py
    # compiled
    python torch_compile_example.py --compile
"""

import argparse
import os
import time

from fastvideo import VideoGenerator

PROMPT = ("A high-definition video of a robotic arm welding a metal structure, "
          "bright sparks and smoke, industrial setting.")


def main() -> None:
    parser = argparse.ArgumentParser(description="torch.compile A/B")
    parser.add_argument("--model", default="Wan-AI/Wan2.1-T2V-1.3B-Diffusers")
    parser.add_argument("--compile", action="store_true", help="Enable torch.compile for the DiT")
    parser.add_argument("--num_gpus", type=int, default=1)
    args = parser.parse_args()

    mode = "COMPILE" if args.compile else "BASELINE"
    print(f"Mode: {mode}  (enable_torch_compile={args.compile})")

    os.makedirs("video_samples", exist_ok=True)

    generator = VideoGenerator.from_pretrained(
        args.model,
        num_gpus=args.num_gpus,
        enable_torch_compile=args.compile,
    )

    def _run(tag: str) -> float:
        save = tag == "measured"
        # Modern typed-request API (generate_video is deprecated). Same
        # prompt/seed/shapes both runs so the compiled graph is reused.
        request: dict = {
            "prompt": PROMPT,
            "sampling": {
                "seed": 1024
            },
            "output": {
                "save_video": save
            },
        }
        if save:
            request["output"]["output_path"] = (f"video_samples/torch_compile_{tag}.mp4")
        t0 = time.perf_counter()
        generator.generate(request)
        return time.perf_counter() - t0

    try:
        # Warmup: pays the one-time graph build when --compile. Discarded.
        w = _run("warmup")
        print(f"warmup: {w:.2f}s "
              f"({'incl. graph build' if args.compile else 'cold start'})")

        # Measured: steady state, compiled graph reused (same shapes/seed).
        m = _run("measured")
        print(f"=== {mode} steady-state e2e: {m:.2f}s ===")
    finally:
        generator.shutdown()


if __name__ == "__main__":
    main()