hao-ai-lab
diff --git a/‎examples/inference/basic/basic_ltx2_3_distilled_i2v.py‎
Lines changed: 298 additions & 0 deletions b/‎examples/inference/basic/basic_ltx2_3_distilled_i2v.py‎
Lines changed: 298 additions & 0 deletions
@@ -0,0 +1,298 @@
+# SPDX-License-Identifier: Apache-2.0
+"""LTX-2.3 distilled image-to-video with per-stage timing breakdown.
+
+This example runs the LTX-2.3 distilled student model on a single GPU and
+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. Runs 2 warmup calls (untimed) + 2 measured calls. The two-warmup pattern
+   is documented because it matters once compile is enabled (see below).
+3. Prints a per-stage breakdown and an average over the measured runs.
+
+torch.compile (opt-in)
+----------------------
+Set `LTX23_ENABLE_COMPILE=1` to enable `torch.compile` on the DiT (with
+fullgraph + max-autotune-no-cudagraphs). With current `main` HEAD and a
+recent PyTorch, this currently fails dynamo at two independent sites — they
+both have to be fixed upstream before the compile path completes:
+
+  1. `fastvideo/models/encoders/gemma.py` (`_replace_padded_with_learnable_
+     registers`): boolean-mask indexing produces a data-dependent shape
+     (`aten.nonzero.default`) that fullgraph + `dynamic=False` can't trace.
+     Workaround: pack valid tokens without a value-dependent intermediate
+     shape (e.g. via `gather`/`scatter` against the binary mask, padding to
+     the original length). Setting `enable_torch_compile_text_encoder=False`
+     here sidesteps it for the text encoder, but the DiT path also trips on
+     it transitively when the connector is compiled — best fixed at the
+     source.
+  2. `fastvideo/models/dits/ltx2.py` around the nested `_build_attn_keep_
+     mask`: the nested-function annotation `bool | torch.Tensor` is
+     evaluated at every outer-forward call, producing a `types.UnionType`
+     that dynamo's SourcelessBuilder cannot wrap. Workaround: hoist the
+     annotation to `typing.Union[bool, torch.Tensor]`, add `from __future__
+     import annotations` at the top of the module, or move the helper to
+     module scope.
+
+Once those two are addressed, this example's compile path should reach
+~4s e2e per clip on GB200 (warm). Until then, it remains eager-by-default.
+
+Hardware notes
+--------------
+- Single-GPU example; for multi-GPU sequence-parallel see the gradio demo
+  under `examples/inference/gradio/local/gradio_local_demo_ltx2_3/`.
+- On GB200 / Blackwell, launch 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 (it is a no-op in eager).
+"""
+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.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)
+ENABLE_COMPILE = os.getenv("LTX23_ENABLE_COMPILE", "0").lower() in (
+    "1", "true", "yes"
+)
+
+# Per-stage timing helpers --------------------------------------------------
+
+def _print_stage_breakdown(result: dict, label: str) -> float | None:
+    """Print stage execution times and return the sum, or None if missing."""
+    logging_info = result.get("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: dict,
+    stage_times: dict[str, list[float]],
+    stage_order: OrderedDict[str, None],
+) -> None:
+    logging_info = result.get("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()}")
+    print(f"Compile:         {'on' if ENABLE_COMPILE else 'off (eager)'}")
+
+    torch_compile_kwargs = {
+        "backend": "inductor",
+        "fullgraph": True,
+        "mode": "max-autotune-no-cudagraphs",
+        "dynamic": False,
+    }
+
+    generator = VideoGenerator.from_pretrained(
+        model_root,
+        num_gpus=1,
+        # LTX-2.3 distilled uses the two-stage refine pipeline; the refine
+        # LoRA is intentionally empty for the distilled student.
+        ltx2_refine_enabled=True,
+        ltx2_refine_upsampler_path=str(refine_upsampler_path),
+        ltx2_refine_lora_path="",
+        ltx2_refine_num_inference_steps=3,
+        ltx2_refine_guidance_scale=1.0,
+        ltx2_refine_add_noise=True,
+        enable_torch_compile=ENABLE_COMPILE,
+        # Text-encoder compile is held off independent of the DiT compile
+        # flag: the current Gemma connector path
+        # (`_replace_padded_with_learnable_registers` in
+        # `fastvideo/models/encoders/gemma.py`) uses boolean-mask indexing
+        # whose output shape depends on the input data, which torch.compile
+        # cannot fullgraph-trace under `dynamic=False`. See the module
+        # docstring for the broader compile-path notes.
+        enable_torch_compile_text_encoder=False,
+        torch_compile_kwargs=torch_compile_kwargs,
+        # Keep everything resident — no CPU offload for serving-style runs.
+        dit_cpu_offload=False,
+        text_encoder_cpu_offload=False,
+        vae_cpu_offload=False,
+        ltx2_vae_tiling=False,
+    )
+
+    common_kwargs = dict(
+        prompt=PROMPT,
+        negative_prompt="",         # distilled is CFG-free; no negative needed
+        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
+        # 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.
+        ltx2_images=[(I2V_IMAGE, 0, 1.0)],
+        ltx2_image_crf=0.0,
+        save_video=True,
+    )
+
+    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_video(
+                output_path=str(OUTPUT_DIR / f"_warmup_{w + 1}.mp4"),
+                seed=7,
+                **common_kwargs,
+            )
+            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_video(
+                output_path=str(out_path),
+                seed=2002 + m,
+                **common_kwargs,
+            )
+            wall = time.perf_counter() - t0
+            e2e = (
+                result.get("e2e_latency")
+                if isinstance(result, dict) 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, dict):
+                _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()