|
| 1 | +# SPDX-License-Identifier: Apache-2.0 |
| 2 | +"""LTX-2.3 distilled text-to-video with the optimized NVFP4 inference stack. |
| 3 | +
|
| 4 | +Runs `FastVideo/LTX-2.3-Distilled-Diffusers` on a single GPU with the full |
| 5 | +validated optimization stack: |
| 6 | +
|
| 7 | +* NVFP4 block-scaled linear layers (per-16 E2M1 weights + E4M3 scale factors), |
| 8 | +* ATTN_QAT_INFER FP4 attention (arch-resolved kernel, receipt logged), |
| 9 | +* torch.compile (fullgraph) over the DiT, text encoder, and VAE, |
| 10 | +* single-stage 8-step distilled sampling at guidance 1.0. |
| 11 | +
|
| 12 | +Quick start |
| 13 | +----------- |
| 14 | + # On GB200-class ARM hosts, unset LD_LIBRARY_PATH (see Hardware notes): |
| 15 | + env -u LD_LIBRARY_PATH python examples/inference/ltx2_3/optimized_nvfp4_t2v.py |
| 16 | +
|
| 17 | + # Optional overrides: |
| 18 | + # export LTX23_MODEL_PATH=/path/to/local/snapshot |
| 19 | + # export LTX23_T2V_PROMPT="a red fox running through fresh snow" |
| 20 | + # export LTX23_OUTPUT_DIR=outputs_video/ltx2_3_nvfp4_t2v |
| 21 | +
|
| 22 | +Hardware notes |
| 23 | +-------------- |
| 24 | +- On GB200 / Blackwell, run with `env -u LD_LIBRARY_PATH ...` to avoid a |
| 25 | + system-cuBLAS / torch-cuBLAS mismatch that fails every GEMM (some ARM |
| 26 | + container images ship an LD_LIBRARY_PATH that breaks torch.compile's |
| 27 | + toolchain discovery). The `_inductor.shape_padding = False` line below |
| 28 | + also avoids a pad_mm landmine on the same generation of cards. |
| 29 | +""" |
| 30 | +from __future__ import annotations |
| 31 | + |
| 32 | +import os |
| 33 | +import time |
| 34 | +from pathlib import Path |
| 35 | + |
| 36 | +import torch._inductor.config as _inductor |
| 37 | + |
| 38 | +from fastvideo import VideoGenerator |
| 39 | +from fastvideo.configs.pipelines.base import PipelineConfig |
| 40 | +from fastvideo.layers.quantization.nvfp4_config import NVFP4Config |
| 41 | +from fastvideo.utils import maybe_download_model |
| 42 | + |
| 43 | +# ATTN_QAT_INFER is the FP4 attention half of the NVFP4 deploy contract. It |
| 44 | +# resolves per arch (CUTLASS SageAttention3-FP4 on sm_120a/sm_121a, FP4 FA4 |
| 45 | +# on sm_100a/sm_103a) and logs a one-line receipt of what actually bound. |
| 46 | +os.environ.setdefault("FASTVIDEO_ATTENTION_BACKEND", "ATTN_QAT_INFER") |
| 47 | +os.environ.setdefault("FASTVIDEO_STAGE_LOGGING", "1") |
| 48 | + |
| 49 | +# Inductor knobs. The first one (shape_padding=False) is mandatory on |
| 50 | +# Blackwell to avoid a cuBLAS INVALID_VALUE crash inside pad_mm. The rest |
| 51 | +# are the same matmul-friendliness flags the sibling LTX-2 examples use. |
| 52 | +_inductor.shape_padding = False |
| 53 | +_inductor.conv_1x1_as_mm = True # treat 1x1 convolutions as matrix muls |
| 54 | +_inductor.coordinate_descent_tuning = True |
| 55 | +_inductor.coordinate_descent_check_all_directions = True |
| 56 | +_inductor.epilogue_fusion = False # do not fuse pointwise ops into matmuls |
| 57 | + |
| 58 | +MODEL_ID = os.path.expandvars( |
| 59 | + os.path.expanduser( |
| 60 | + os.getenv("LTX23_MODEL_PATH", "FastVideo/LTX-2.3-Distilled-Diffusers") |
| 61 | + ) |
| 62 | +) |
| 63 | +OUTPUT_DIR = Path(os.getenv("LTX23_OUTPUT_DIR", "outputs_video/ltx2_3_nvfp4_t2v")) |
| 64 | +DEFAULT_PROMPT = ( |
| 65 | + "A fashion model takes a slow step forward and shifts her weight, " |
| 66 | + "the soft fabric of her clothing swaying and rippling with the " |
| 67 | + "motion, her hair shifting gently, soft even studio lighting on a " |
| 68 | + "clean light background, elegant slow-motion runway feel." |
| 69 | +) |
| 70 | +PROMPT = os.getenv("LTX23_T2V_PROMPT", DEFAULT_PROMPT) |
| 71 | + |
| 72 | + |
| 73 | +def main() -> None: |
| 74 | + OUTPUT_DIR.mkdir(parents=True, exist_ok=True) |
| 75 | + model_root = maybe_download_model(MODEL_ID) |
| 76 | + print(f"Model: {model_root}") |
| 77 | + print(f"Output dir: {OUTPUT_DIR.resolve()}") |
| 78 | + |
| 79 | + # Loading the pipeline config *with model_path* binds model-specific |
| 80 | + # tuning (notably VAE precision/decoder defaults) into the config. |
| 81 | + pipeline_config = PipelineConfig.from_pretrained(model_root) |
| 82 | + |
| 83 | + # NVFP4 linear layers on the DiT. The default purges the original BF16 |
| 84 | + # weights of the always-FP4 linears right after conversion — a large |
| 85 | + # peak-memory reduction. Refine-only layers (the cross-modal AV |
| 86 | + # projections) always keep theirs: the base stage profile runs them |
| 87 | + # dense by deployment contract. retain_original_weights=True keeps |
| 88 | + # everything (debugging). |
| 89 | + pipeline_config.dit_config.quant_config = NVFP4Config() |
| 90 | + |
| 91 | + # fullgraph=True is supported on the NVFP4 path: the FP4 quantize step |
| 92 | + # is a registered custom op (fastvideo::nvfp4_quantize_fa4), so dynamo |
| 93 | + # traces through it without graph breaks. mode="default" — the two |
| 94 | + # CUDAGraph modes ("reduce-overhead" / "max-autotune") are a separate |
| 95 | + # opt-in, not part of this validated preset. |
| 96 | + torch_compile_kwargs = { |
| 97 | + "backend": "inductor", |
| 98 | + "fullgraph": True, |
| 99 | + "mode": "default", |
| 100 | + "dynamic": False, |
| 101 | + } |
| 102 | + |
| 103 | + generator = VideoGenerator.from_pretrained( |
| 104 | + model_root, |
| 105 | + num_gpus=1, |
| 106 | + pipeline_config=pipeline_config, |
| 107 | + # Compile the DiT, text encoder, and VAE — all three stages benefit, |
| 108 | + # and the VAE's codec submodules compile cleanly under fullgraph. |
| 109 | + enable_torch_compile=True, |
| 110 | + enable_torch_compile_text_encoder=True, |
| 111 | + enable_torch_compile_vae=True, |
| 112 | + torch_compile_kwargs=torch_compile_kwargs, |
| 113 | + torch_compile_kwargs_vae=torch_compile_kwargs, |
| 114 | + # Keep everything resident — no CPU offload for serving-style runs. |
| 115 | + dit_cpu_offload=False, |
| 116 | + text_encoder_cpu_offload=False, |
| 117 | + vae_cpu_offload=False, |
| 118 | + ltx2_vae_tiling=False, |
| 119 | + ) |
| 120 | + |
| 121 | + common_kwargs = dict( |
| 122 | + prompt=PROMPT, |
| 123 | + negative_prompt="", # distilled is CFG-free; no negative needed |
| 124 | + guidance_scale=1.0, # CFG=1 for distilled |
| 125 | + height=1280, width=832, # portrait runway aspect |
| 126 | + num_frames=121, fps=24, # ~5s clip |
| 127 | + # Single-stage 8-step distilled sampling — the validated preset for |
| 128 | + # this checkpoint (no two-stage refine; the NVFP4 deploy contract |
| 129 | + # runs the distilled single-stage recipe). |
| 130 | + num_inference_steps=8, |
| 131 | + save_video=True, |
| 132 | + ) |
| 133 | + |
| 134 | + try: |
| 135 | + # Warmup: pays cold compile + first-shape guard work, untimed. |
| 136 | + print("\n[warmup] compiling + generating…") |
| 137 | + generator.generate_video( |
| 138 | + output_path=str(OUTPUT_DIR / "_warmup.mp4"), |
| 139 | + seed=7, |
| 140 | + **common_kwargs, |
| 141 | + ) |
| 142 | + (OUTPUT_DIR / "_warmup.mp4").unlink(missing_ok=True) |
| 143 | + |
| 144 | + # Measured run. |
| 145 | + out_path = OUTPUT_DIR / "output_ltx2_3_nvfp4_t2v.mp4" |
| 146 | + print(f"\n[measured] generating: {out_path}") |
| 147 | + t0 = time.perf_counter() |
| 148 | + result = generator.generate_video( |
| 149 | + output_path=str(out_path), |
| 150 | + seed=2002, |
| 151 | + **common_kwargs, |
| 152 | + ) |
| 153 | + wall = time.perf_counter() - t0 |
| 154 | + e2e = (result.get("e2e_latency") if isinstance(result, dict) else None) or wall |
| 155 | + print(f"[measured] e2e={e2e:.2f}s wall={wall:.2f}s") |
| 156 | + finally: |
| 157 | + generator.shutdown() |
| 158 | + |
| 159 | + # Expected receipts — verify these two lines in your own run's log: |
| 160 | + # |
| 161 | + # 1. ATTN_QAT_INFER routing receipt (logged at backend resolution; on a |
| 162 | + # GB200-class part it reads): |
| 163 | + # |
| 164 | + # ATTN_QAT_INFER resolved: arch=sm_100 kernel=flash-attention-fp4 \ |
| 165 | + # qk_mode=nvfp4(per-16-e4m3-sf) pv_mode=bf16 train_sim_mismatch=measured |
| 166 | + # |
| 167 | + # 2. NVFP4 weight purge receipt (logged after model conversion; N/M/X |
| 168 | + # depend on the checkpoint): |
| 169 | + # |
| 170 | + # NVFP4 weight purge receipt: purged N original bf16 weight tensors \ |
| 171 | + # (X.XX GiB freed); retained M (refine-only dense fallback or \ |
| 172 | + # retain_original_weights). |
| 173 | + |
| 174 | + |
| 175 | +if __name__ == "__main__": |
| 176 | + main() |
0 commit comments