Files
hao-ai-lab-FastVideo/examples/inference/basic/basic_ltx2_3_distilled_i2v.py
T
Davids048andClaude Opus 5.5 f571621ae7 [refactor]: move examples to generate(request) and from_config
Every example calls VideoGenerator.generate with a nested request instead
of generate_video keywords, and passes settings outside the from_pretrained
convenience keywords through VideoGenerator.from_config at their typed
paths. Each request and generator config equals what the keyword call built.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_012M91wnVFmPEvJ9h39r5BH7
2026-10-03 02:14:36 +00:00

293 lines
12 KiB
Python

# SPDX-License-Identifier: Apache-2.0
"""LTX-2.3 distilled image-to-video with torch.compile + timing breakdown.
This example runs the LTX-2.3 distilled student model on a single GPU with
torch.compile fully enabled, then prints a per-stage timing breakdown so the
user can see where wall-time goes. It is meant as the canonical entry point
for trying out the LTX-2.3 i2v path on `hao-ai-lab/FastVideo:main`.
Quick start
-----------
export LTX23_I2V_IMAGE=/path/to/your/portrait_or_product.jpg
# optional overrides:
# export LTX23_I2V_PROMPT="a fashion model walks toward camera..."
# export LTX23_OUTPUT_DIR=outputs_video/ltx2_3_distilled_i2v
python examples/inference/basic/basic_ltx2_3_distilled_i2v.py
What the script does
--------------------
1. Loads FastVideo/LTX-2.3-Distilled-Diffusers (8 denoise + 3 refine steps,
CFG=1, no refine LoRA — the distilled production recipe).
2. Compiles the DiT, text encoder, and VAE (fullgraph, Inductor default
mode — autotune adds ~7 min cold-compile here with no measurable
e2e gain).
3. Runs 2 warmup calls (untimed) + 2 measured calls. Two warmups are kept
as a safety net — the first call pays cold compile + first-shape guard
work, and a second warmup ensures any residual recompiles settle before
we measure.
4. Prints a per-stage breakdown and an average over the measured runs.
Hardware notes
--------------
- Single-GPU example; for multi-GPU sequence-parallel see the gradio demo
under `examples/inference/gradio/local/gradio_local_demo_ltx2_3/`.
- First-time compile takes ~30-40 min on GB200 (~20 min on H100; cached
in `$TORCHINDUCTOR_CACHE_DIR` afterwards). Subsequent invocations only
pay the one-time process load + a few seconds of dynamo trace.
- On GB200 / Blackwell, run with `env -u LD_LIBRARY_PATH ...` to avoid a
system-cuBLAS / torch-cuBLAS mismatch that fails every GEMM. The
`_inductor.shape_padding = False` line below also avoids a pad_mm
landmine on the same generation of cards.
"""
from __future__ import annotations
import os
import time
from collections import OrderedDict
from pathlib import Path
import torch._inductor.config as _inductor
from fastvideo import VideoGenerator
from fastvideo.api import GenerationResult
from fastvideo.configs.pipelines.base import PipelineConfig
from fastvideo.utils import maybe_download_model
# Env knobs (set BEFORE importing fastvideo where possible — but
# FASTVIDEO_ATTENTION_BACKEND is fine here because the worker reads it
# on generator construction).
os.environ.setdefault("FASTVIDEO_ATTENTION_BACKEND", "FLASH_ATTN")
os.environ.setdefault("FASTVIDEO_STAGE_LOGGING", "1")
# Inductor knobs. The first one (shape_padding=False) is mandatory on
# Blackwell to avoid a cuBLAS INVALID_VALUE crash inside pad_mm during
# the refine path. The rest are autotune-friendliness flags.
_inductor.shape_padding = False
_inductor.conv_1x1_as_mm = True
_inductor.coordinate_descent_tuning = True
_inductor.coordinate_descent_check_all_directions = True
_inductor.epilogue_fusion = False
MODEL_ID = os.path.expandvars(os.path.expanduser(os.getenv("LTX23_MODEL_PATH",
"FastVideo/LTX-2.3-Distilled-Diffusers")))
OUTPUT_DIR = Path(os.getenv("LTX23_OUTPUT_DIR", "outputs_video/ltx2_3_distilled_i2v"))
I2V_IMAGE = os.getenv("LTX23_I2V_IMAGE", "")
DEFAULT_PROMPT = ("A fashion model takes a slow step forward and shifts her weight, "
"the soft fabric of her clothing swaying and rippling with the "
"motion, her hair shifting gently, soft even studio lighting on a "
"clean light background, elegant slow-motion runway feel.")
PROMPT = os.getenv("LTX23_I2V_PROMPT", DEFAULT_PROMPT)
# Per-stage timing helpers --------------------------------------------------
def _print_stage_breakdown(result: GenerationResult, label: str) -> float | None:
"""Print stage execution times and return the sum, or None if missing."""
logging_info = result.logging_info
stages = getattr(logging_info, "stages", None) if logging_info else None
if not stages:
print(f" [{label}] stage breakdown unavailable")
return None
print(f" [{label}] stage breakdown:")
total = 0.0
for name, metrics in stages.items():
exec_s = float(metrics.get("execution_time", 0.0))
total += exec_s
print(f" - {name}: {exec_s:.3f}s")
print(f" - stage_sum: {total:.3f}s")
return total
def _collect_stage_times(
result: GenerationResult,
stage_times: dict[str, list[float]],
stage_order: OrderedDict[str, None],
) -> None:
logging_info = result.logging_info
stages = getattr(logging_info, "stages", None) if logging_info else None
if not stages:
return
for name, metrics in stages.items():
stage_order.setdefault(name, None)
stage_times.setdefault(name, []).append(float(metrics.get("execution_time", 0.0)))
def _resolve_refine_upsampler(model_root: str) -> Path:
"""LTX-2.3 distilled snapshots ship a `spatial_upscaler/` subdir."""
for name in ("spatial_upscaler", "spatial_upsampler"):
cand = Path(model_root) / name
if (cand / "config.json").is_file():
return cand
raise FileNotFoundError(f"No refine upsampler directory under {model_root}. "
f"Expected `{model_root}/spatial_upscaler/config.json`.")
# Main ---------------------------------------------------------------------
def main() -> None:
if not I2V_IMAGE:
raise SystemExit("LTX23_I2V_IMAGE is required for i2v. Example:\n"
" export LTX23_I2V_IMAGE=/path/to/portrait_or_product.jpg\n"
" python examples/inference/basic/basic_ltx2_3_distilled_i2v.py")
if not Path(I2V_IMAGE).is_file():
raise SystemExit(f"LTX23_I2V_IMAGE not found: {I2V_IMAGE}")
OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
model_root = maybe_download_model(MODEL_ID)
refine_upsampler_path = _resolve_refine_upsampler(model_root)
print(f"Model: {model_root}")
print(f"Refine upsampler: {refine_upsampler_path}")
print(f"i2v image: {I2V_IMAGE}")
print(f"Output dir: {OUTPUT_DIR.resolve()}")
# mode="default" — Inductor's default schedule matches max-autotune on
# this pipeline (denoise/refine/decode all within ~5 ms, n=2) while
# saving ~7 min of cold compile on a single GB200.
torch_compile_kwargs = {
"backend": "inductor",
"fullgraph": True,
"mode": "default",
"dynamic": False,
}
# Loading the pipeline config *with model_path* binds model-specific
# tuning (notably VAE precision/decoder defaults) into the config. Without
# this, the generic pipeline config gives a substantially slower VAE
# decode stage. `basic_ltx2_distilled_fast_profile.py` uses the same
# pattern.
pipeline_config = PipelineConfig.from_pretrained(model_root)
pipeline_config.dit_config.quant_config = None
generator = VideoGenerator.from_config({
"model_path": model_root,
"engine": {
"num_gpus": 1,
"compile": {
"enabled": True,
"text_encoder_enabled": True,
# Compile the VAE codec submodules (encoder / decoder) too. The
# `LTX2CausalVideoAutoencoder` declares `_compile_conditions` so
# `_compile_with_conditions` targets just those submodules and
# leaves the surrounding tiling control flow eager — needed for
# fullgraph + dynamic=False to succeed. VAE eager decode is
# ~1.0s; compiling it brings the stage to ~0.3s.
"vae_enabled": True,
**torch_compile_kwargs,
"vae_kwargs": torch_compile_kwargs,
},
# Keep everything resident — no CPU offload for serving-style runs.
"offload": {
"dit": False,
"text_encoder": False,
"vae": False,
},
},
"pipeline": {
# LTX-2.3 distilled uses the two-stage refine pipeline; the refine
# LoRA is intentionally empty for the distilled student.
"preset_overrides": {
"refine": {
"enabled": True,
"num_inference_steps": 3,
"guidance_scale": 1.0,
"add_noise": True,
},
},
"components": {"upsampler_weights": str(refine_upsampler_path)},
"ltx2": {"refine": {"lora_path": ""}},
"experimental": {"pipeline_config": pipeline_config},
"vae_tiling": False,
},
})
common_request = {
"prompt": PROMPT,
"negative_prompt": "", # distilled is CFG-free; no negative needed
"sampling": {
"guidance_scale": 1.0, # CFG=1 for distilled
"height": 1280,
"width": 832, # portrait runway aspect
"num_frames": 121,
"fps": 24, # ~5s clip
"num_inference_steps": 8, # distilled denoise steps
},
"output": {"save_video": True},
# i2v: anchor the input image at frame 0 with full strength.
# `ltx2_image_crf=0.0` skips an extra JPEG re-encode of an already
# JPEG conditioning image.
"extensions": {
"ltx2_images": [(I2V_IMAGE, 0, 1.0)],
"ltx2_image_crf": 0.0,
},
}
warmup_runs = 2
measured_runs = 2
warmup_secs: list[float] = []
measured_secs: list[float] = []
stage_times: dict[str, list[float]] = {}
stage_order: OrderedDict[str, None] = OrderedDict()
try:
# Warmup: untimed (but we still wall-clock them so the first compile
# cost is visible to the reader).
for w in range(warmup_runs):
t0 = time.perf_counter()
print(f"\n[warmup {w + 1}/{warmup_runs}] compiling + generating…")
generator.generate({
**common_request,
"sampling": {**common_request["sampling"], "seed": 7},
"output": {**common_request["output"], "output_path": str(OUTPUT_DIR / f"_warmup_{w + 1}.mp4")},
})
dt = time.perf_counter() - t0
warmup_secs.append(dt)
print(f"[warmup {w + 1}/{warmup_runs}] wall={dt:.1f}s")
# Cleanup warmup artifacts so the user only sees measured outputs.
for w in range(warmup_runs):
(OUTPUT_DIR / f"_warmup_{w + 1}.mp4").unlink(missing_ok=True)
# Measured.
for m in range(measured_runs):
out_path = OUTPUT_DIR / f"output_ltx2_3_distilled_i2v_run_{m + 1}.mp4"
print(f"\n[measured {m + 1}/{measured_runs}] generating: {out_path}")
t0 = time.perf_counter()
result = generator.generate({
**common_request,
"sampling": {**common_request["sampling"], "seed": 2002 + m},
"output": {**common_request["output"], "output_path": str(out_path)},
})
wall = time.perf_counter() - t0
e2e = (result.extra.get("e2e_latency") if isinstance(result, GenerationResult) else None) or wall
measured_secs.append(e2e)
print(f"[measured {m + 1}/{measured_runs}] e2e={e2e:.2f}s wall={wall:.2f}s")
if isinstance(result, GenerationResult):
_print_stage_breakdown(result, f"measured {m + 1}")
_collect_stage_times(result, stage_times, stage_order)
# Summary.
print("\n=== summary ===")
print(f"warmup wall-times: {[round(x, 1) for x in warmup_secs]}")
if measured_secs:
avg = sum(measured_secs) / len(measured_secs)
print(f"measured e2e (n={len(measured_secs)}): "
f"{[round(x, 2) for x in measured_secs]} -> avg {avg:.2f}s")
if stage_times:
print(f"average stage times over {measured_runs} measured runs:")
avg_total = 0.0
for name in stage_order:
vals = stage_times.get(name) or []
if not vals:
continue
avg_v = sum(vals) / len(vals)
avg_total += avg_v
print(f" - {name}: {avg_v:.3f}s")
print(f" - stage_sum_avg: {avg_total:.3f}s")
finally:
generator.shutdown()
if __name__ == "__main__":
main()