Skip to content

Commit a50bd67

Browse files
committed
[example] LTX-2.3 distilled i2v with compile + timing breakdown
Add `examples/inference/basic/basic_ltx2_3_distilled_i2v.py`: a single-GPU LTX-2.3 distilled image-to-video example with torch.compile fully enabled, two warmup runs to settle Inductor's per-shape autotune, two measured runs, and a per-stage timing breakdown. Fills a gap on `main`: the existing `basic_ltx2*.py` examples are t2v-only and target LTX-2.0 distilled. After PR #1397 merged LTX-2.3, there was no copy-paste example for the LTX-2.3 i2v path with compile + benchmark plumbing wired in. The script reads the conditioning image from `LTX23_I2V_IMAGE` (errors out with a helpful message if unset) and uses a generic fashion-runway prompt that the user can override via `LTX23_I2V_PROMPT`. Defaults match the production recipe documented in the docstring: 8 denoise + 3 refine steps, CFG=1, 832x1280 portrait, 121 frames @ 24fps. Includes a comment calling out the Blackwell `shape_padding=False` requirement and the `env -u LD_LIBRARY_PATH` launch tip from prior experience on GB200.
1 parent 5706079 commit a50bd67

1 file changed

Lines changed: 268 additions & 0 deletions

File tree

Lines changed: 268 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,268 @@
1+
# SPDX-License-Identifier: Apache-2.0
2+
"""LTX-2.3 distilled image-to-video with torch.compile + timing breakdown.
3+
4+
This example runs the LTX-2.3 distilled student model on a single GPU with
5+
torch.compile fully enabled, then prints a per-stage timing breakdown so the
6+
user can see where wall-time goes. It is meant as the canonical entry point
7+
for trying out the LTX-2.3 i2v path on `hao-ai-lab/FastVideo:main`.
8+
9+
Quick start
10+
-----------
11+
export LTX23_I2V_IMAGE=/path/to/your/portrait_or_product.jpg
12+
# optional overrides:
13+
# export LTX23_I2V_PROMPT="a fashion model walks toward camera..."
14+
# export LTX23_OUTPUT_DIR=outputs_video/ltx2_3_distilled_i2v
15+
python examples/inference/basic/basic_ltx2_3_distilled_i2v.py
16+
17+
What the script does
18+
--------------------
19+
1. Loads FastVideo/LTX-2.3-Distilled-Diffusers (8 denoise + 3 refine steps,
20+
CFG=1, no refine LoRA — the distilled production recipe).
21+
2. Compiles the DiT, text encoder, and VAE (fullgraph, max-autotune-no-
22+
cudagraphs).
23+
3. Runs 2 warmup calls (untimed) + 2 measured calls. Two warmups are needed
24+
even though distilled has no refine LoRA, because Inductor's per-shape
25+
autotune for the first call still leaves a few cold guards on call 2;
26+
the second warmup settles them. Skipping to a single warmup typically
27+
inflates the first measured run by tens of seconds.
28+
4. Prints a per-stage breakdown and an average over the measured runs.
29+
30+
Hardware notes
31+
--------------
32+
- Single-GPU example; for multi-GPU sequence-parallel see the gradio demo
33+
under `examples/inference/gradio/local/gradio_local_demo_ltx2_3/`.
34+
- First-time compile + autotune takes ~10-30 min on H100 / GB200 (cached
35+
in `$TORCHINDUCTOR_CACHE_DIR` afterwards). Subsequent invocations only
36+
pay the one-time process load + a few seconds of dynamo trace.
37+
- On GB200 / Blackwell, run with `env -u LD_LIBRARY_PATH ...` to avoid a
38+
system-cuBLAS / torch-cuBLAS mismatch that fails every GEMM. The
39+
`_inductor.shape_padding = False` line below also avoids a pad_mm
40+
landmine on the same generation of cards.
41+
"""
42+
from __future__ import annotations
43+
44+
import os
45+
import time
46+
from collections import OrderedDict
47+
from pathlib import Path
48+
49+
import torch._inductor.config as _inductor
50+
51+
from fastvideo import VideoGenerator
52+
from fastvideo.utils import maybe_download_model
53+
54+
# Env knobs (set BEFORE importing fastvideo where possible — but
55+
# FASTVIDEO_ATTENTION_BACKEND is fine here because the worker reads it
56+
# on generator construction).
57+
os.environ.setdefault("FASTVIDEO_ATTENTION_BACKEND", "FLASH_ATTN")
58+
os.environ.setdefault("FASTVIDEO_STAGE_LOGGING", "1")
59+
60+
# Inductor knobs. The first one (shape_padding=False) is mandatory on
61+
# Blackwell to avoid a cuBLAS INVALID_VALUE crash inside pad_mm during
62+
# the refine path. The rest are autotune-friendliness flags.
63+
_inductor.shape_padding = False
64+
_inductor.conv_1x1_as_mm = True
65+
_inductor.coordinate_descent_tuning = True
66+
_inductor.coordinate_descent_check_all_directions = True
67+
_inductor.epilogue_fusion = False
68+
69+
MODEL_ID = os.path.expandvars(
70+
os.path.expanduser(
71+
os.getenv("LTX23_MODEL_PATH", "FastVideo/LTX-2.3-Distilled-Diffusers")
72+
)
73+
)
74+
OUTPUT_DIR = Path(
75+
os.getenv("LTX23_OUTPUT_DIR", "outputs_video/ltx2_3_distilled_i2v")
76+
)
77+
I2V_IMAGE = os.getenv("LTX23_I2V_IMAGE", "")
78+
DEFAULT_PROMPT = (
79+
"A fashion model takes a slow step forward and shifts her weight, "
80+
"the soft fabric of her clothing swaying and rippling with the "
81+
"motion, her hair shifting gently, soft even studio lighting on a "
82+
"clean light background, elegant slow-motion runway feel."
83+
)
84+
PROMPT = os.getenv("LTX23_I2V_PROMPT", DEFAULT_PROMPT)
85+
86+
# Per-stage timing helpers --------------------------------------------------
87+
88+
def _print_stage_breakdown(result: dict, label: str) -> float | None:
89+
"""Print stage execution times and return the sum, or None if missing."""
90+
logging_info = result.get("logging_info")
91+
stages = getattr(logging_info, "stages", None) if logging_info else None
92+
if not stages:
93+
print(f" [{label}] stage breakdown unavailable")
94+
return None
95+
print(f" [{label}] stage breakdown:")
96+
total = 0.0
97+
for name, metrics in stages.items():
98+
exec_s = float(metrics.get("execution_time", 0.0))
99+
total += exec_s
100+
print(f" - {name}: {exec_s:.3f}s")
101+
print(f" - stage_sum: {total:.3f}s")
102+
return total
103+
104+
105+
def _collect_stage_times(
106+
result: dict,
107+
stage_times: dict[str, list[float]],
108+
stage_order: OrderedDict[str, None],
109+
) -> None:
110+
logging_info = result.get("logging_info")
111+
stages = getattr(logging_info, "stages", None) if logging_info else None
112+
if not stages:
113+
return
114+
for name, metrics in stages.items():
115+
stage_order.setdefault(name, None)
116+
stage_times.setdefault(name, []).append(
117+
float(metrics.get("execution_time", 0.0))
118+
)
119+
120+
121+
def _resolve_refine_upsampler(model_root: str) -> Path:
122+
"""LTX-2.3 distilled snapshots ship a `spatial_upscaler/` subdir."""
123+
for name in ("spatial_upscaler", "spatial_upsampler"):
124+
cand = Path(model_root) / name
125+
if (cand / "config.json").is_file():
126+
return cand
127+
raise FileNotFoundError(
128+
f"No refine upsampler directory under {model_root}. "
129+
f"Expected `{model_root}/spatial_upscaler/config.json`."
130+
)
131+
132+
133+
# Main ---------------------------------------------------------------------
134+
135+
def main() -> None:
136+
if not I2V_IMAGE:
137+
raise SystemExit(
138+
"LTX23_I2V_IMAGE is required for i2v. Example:\n"
139+
" export LTX23_I2V_IMAGE=/path/to/portrait_or_product.jpg\n"
140+
" python examples/inference/basic/basic_ltx2_3_distilled_i2v.py"
141+
)
142+
if not Path(I2V_IMAGE).is_file():
143+
raise SystemExit(f"LTX23_I2V_IMAGE not found: {I2V_IMAGE}")
144+
145+
OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
146+
model_root = maybe_download_model(MODEL_ID)
147+
refine_upsampler_path = _resolve_refine_upsampler(model_root)
148+
print(f"Model: {model_root}")
149+
print(f"Refine upsampler: {refine_upsampler_path}")
150+
print(f"i2v image: {I2V_IMAGE}")
151+
print(f"Output dir: {OUTPUT_DIR.resolve()}")
152+
153+
torch_compile_kwargs = {
154+
"backend": "inductor",
155+
"fullgraph": True,
156+
"mode": "max-autotune-no-cudagraphs",
157+
"dynamic": False,
158+
}
159+
160+
generator = VideoGenerator.from_pretrained(
161+
model_root,
162+
num_gpus=1,
163+
# LTX-2.3 distilled uses the two-stage refine pipeline; the refine
164+
# LoRA is intentionally empty for the distilled student.
165+
ltx2_refine_enabled=True,
166+
ltx2_refine_upsampler_path=str(refine_upsampler_path),
167+
ltx2_refine_lora_path="",
168+
ltx2_refine_num_inference_steps=3,
169+
ltx2_refine_guidance_scale=1.0,
170+
ltx2_refine_add_noise=True,
171+
enable_torch_compile=True,
172+
enable_torch_compile_text_encoder=True,
173+
torch_compile_kwargs=torch_compile_kwargs,
174+
# Keep everything resident — no CPU offload for serving-style runs.
175+
dit_cpu_offload=False,
176+
text_encoder_cpu_offload=False,
177+
vae_cpu_offload=False,
178+
ltx2_vae_tiling=False,
179+
)
180+
181+
common_kwargs = dict(
182+
prompt=PROMPT,
183+
negative_prompt="", # distilled is CFG-free; no negative needed
184+
guidance_scale=1.0, # CFG=1 for distilled
185+
height=1280, width=832, # portrait runway aspect
186+
num_frames=121, fps=24, # ~5s clip
187+
num_inference_steps=8, # distilled denoise steps
188+
# i2v: anchor the input image at frame 0 with full strength.
189+
# `ltx2_image_crf=0.0` skips an extra JPEG re-encode of an already
190+
# JPEG conditioning image.
191+
ltx2_images=[(I2V_IMAGE, 0, 1.0)],
192+
ltx2_image_crf=0.0,
193+
save_video=True,
194+
)
195+
196+
warmup_runs = 2
197+
measured_runs = 2
198+
warmup_secs: list[float] = []
199+
measured_secs: list[float] = []
200+
stage_times: dict[str, list[float]] = {}
201+
stage_order: OrderedDict[str, None] = OrderedDict()
202+
203+
try:
204+
# Warmup: untimed (but we still wall-clock them so the first compile
205+
# cost is visible to the reader).
206+
for w in range(warmup_runs):
207+
t0 = time.perf_counter()
208+
print(f"\n[warmup {w + 1}/{warmup_runs}] compiling + generating…")
209+
generator.generate_video(
210+
output_path=str(OUTPUT_DIR / f"_warmup_{w + 1}.mp4"),
211+
seed=7,
212+
**common_kwargs,
213+
)
214+
dt = time.perf_counter() - t0
215+
warmup_secs.append(dt)
216+
print(f"[warmup {w + 1}/{warmup_runs}] wall={dt:.1f}s")
217+
218+
# Cleanup warmup artifacts so the user only sees measured outputs.
219+
for w in range(warmup_runs):
220+
(OUTPUT_DIR / f"_warmup_{w + 1}.mp4").unlink(missing_ok=True)
221+
222+
# Measured.
223+
for m in range(measured_runs):
224+
out_path = OUTPUT_DIR / f"output_ltx2_3_distilled_i2v_run_{m + 1}.mp4"
225+
print(f"\n[measured {m + 1}/{measured_runs}] generating: {out_path}")
226+
t0 = time.perf_counter()
227+
result = generator.generate_video(
228+
output_path=str(out_path),
229+
seed=2002 + m,
230+
**common_kwargs,
231+
)
232+
wall = time.perf_counter() - t0
233+
e2e = (
234+
result.get("e2e_latency")
235+
if isinstance(result, dict) else None
236+
) or wall
237+
measured_secs.append(e2e)
238+
print(f"[measured {m + 1}/{measured_runs}] e2e={e2e:.2f}s wall={wall:.2f}s")
239+
if isinstance(result, dict):
240+
_print_stage_breakdown(result, f"measured {m + 1}")
241+
_collect_stage_times(result, stage_times, stage_order)
242+
243+
# Summary.
244+
print("\n=== summary ===")
245+
print(f"warmup wall-times: {[round(x, 1) for x in warmup_secs]}")
246+
if measured_secs:
247+
avg = sum(measured_secs) / len(measured_secs)
248+
print(
249+
f"measured e2e (n={len(measured_secs)}): "
250+
f"{[round(x, 2) for x in measured_secs]} -> avg {avg:.2f}s"
251+
)
252+
if stage_times:
253+
print(f"average stage times over {measured_runs} measured runs:")
254+
avg_total = 0.0
255+
for name in stage_order:
256+
vals = stage_times.get(name) or []
257+
if not vals:
258+
continue
259+
avg_v = sum(vals) / len(vals)
260+
avg_total += avg_v
261+
print(f" - {name}: {avg_v:.3f}s")
262+
print(f" - stage_sum_avg: {avg_total:.3f}s")
263+
finally:
264+
generator.shutdown()
265+
266+
267+
if __name__ == "__main__":
268+
main()

0 commit comments

Comments
 (0)