Upload demo/demo_distilled_decode.py

demo/demo_distilled_decode.py

@@ -0,0 +1,219 @@
+#!/usr/bin/env python3
+"""Quick demo: generate a short LTX-2.3 video latent, then decode it twice.
+
+Compares the stock LTX-2.3 VAE decoder with this repo's pruned decoder
+(PrunaVAED) on the *same* latent. Prints decode time (ms) and writes two mp4s.
+
+Requires a CUDA GPU. From the repo root:
+
+    pip install -r requirements-test.txt
+    python tests/test_distilled_decode.py
+"""
+
+from __future__ import annotations
+
+import statistics
+import sys
+import time
+from pathlib import Path
+
+import imageio.v3 as iio
+import torch
+from diffusers import LTX2LatentUpsamplePipeline, LTX2Pipeline
+from diffusers.models.autoencoders import AutoencoderKLLTX2Video
+from diffusers.pipelines.ltx2.latent_upsampler import LTX2LatentUpsamplerModel
+from diffusers.pipelines.ltx2.utils import (
+    DEFAULT_NEGATIVE_PROMPT,
+    DISTILLED_SIGMA_VALUES,
+    STAGE_2_DISTILLED_SIGMA_VALUES,
+)
+
+torch.backends.cuda.enable_cudnn_sdp(False)
+
+# ---------------------------------------------------------------------------
+# Settings (edit these if you want)
+# ---------------------------------------------------------------------------
+
+REPO_ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(REPO_ROOT))
+
+from patch_diffusers import patch_pruna_ltx2_decoder  # noqa: E402
+
+# Official Diffusers checkpoints for the 2-stage distilled recipe.
+DISTILLED_MODEL = "diffusers/LTX-2.3-Distilled-Diffusers"
+SPATIAL_UPSAMPLER = "dg845/LTX-2.3-Spatial-Upsampler-Diffusers"
+LTX23_VAE = "diffusers/LTX-2.3-Diffusers"  # stock decoder (baseline)
+PRUNED_VAE = REPO_ROOT  # this repo's vae/ folder
+
+# ~720p for 5 s @ 24 fps. Height/width must be multiples of 64 (2-stage).
+HEIGHT, WIDTH, NUM_FRAMES, FPS = 1088, 1920, 121, 24.0
+SEED = 42
+DECODE_WARMUP, DECODE_RUNS = 1, 3  # timing: 1 warm-up + median of 3 runs
+
+PROMPT = (
+"The video shows a hockey player in a green jersey and blue helmet skating on the ice with a hockey stick. The player is seen moving around the ice, passing the puck to another player who is also wearing a green jersey and blue helmet. The player in the green jersey is seen skating away from the camera, and then turning around to face the camera. The ice rink is surrounded by boards with advertisements, and there are other players in the background. The player in the green jersey is wearing black gloves and black skates. The player in the green jersey is also seen skating towards the camera and then away from the camera again."
+)
+
+DEVICE = "cuda"
+DTYPE = torch.bfloat16
+OUTPUT_DIR = REPO_ROOT / "outputs" / "test_distilled"
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+def save_mp4(video: torch.Tensor, path: Path) -> None:
+    """Save a BCTHW tensor in [-1, 1] as an H.264 mp4."""
+    frames = video[0].permute(1, 2, 3, 0).clamp(-1, 1)
+    frames = ((frames + 1) / 2 * 255).byte().cpu().numpy()
+    path.parent.mkdir(parents=True, exist_ok=True)
+    iio.imwrite(path, frames, fps=FPS, codec="libx264")
+
+
+def load_vae(model_id: str) -> AutoencoderKLLTX2Video:
+    """Load a video VAE decoder and put it on GPU."""
+    vae = AutoencoderKLLTX2Video.from_pretrained(
+        model_id, subfolder="vae", torch_dtype=DTYPE
+    )
+    vae = vae.to(DEVICE).eval()
+    vae.use_tiling = False
+    if hasattr(vae, "disable_tiling"):
+        vae.disable_tiling()
+    return vae
+
+
+@torch.inference_mode()
+def timed_decode(vae: AutoencoderKLLTX2Video, latent: torch.Tensor) -> tuple[torch.Tensor, float]:
+    """Decode once after warm-up; return (video on CPU, median latency in ms)."""
+    latent = latent.to(DEVICE, DTYPE)
+
+    with torch.autocast(DEVICE, dtype=DTYPE):
+        for _ in range(DECODE_WARMUP):
+            vae.decode(latent, return_dict=False)
+
+    times_ms = []
+    video = None
+    for _ in range(DECODE_RUNS):
+        torch.cuda.synchronize()
+        t0 = time.perf_counter()
+        with torch.autocast(DEVICE, dtype=DTYPE):
+            video = vae.decode(latent, return_dict=False)[0]
+        torch.cuda.synchronize()
+        times_ms.append((time.perf_counter() - t0) * 1000)
+
+    return video.cpu(), statistics.median(times_ms)
+
+
+@torch.inference_mode()
+def generate_latent(prompt: str) -> torch.Tensor:
+    """Run the official 2-stage distilled pipeline and return the final video latent.
+
+    Same layout as Lightricks DistilledPipeline:
+      stage 1 @ half-res → 2× latent upsample → stage 2 @ full-res.
+    """
+    half_h, half_w = HEIGHT // 2, WIDTH // 2
+    generator = torch.Generator(DEVICE).manual_seed(SEED)
+
+    pipe = LTX2Pipeline.from_pretrained(DISTILLED_MODEL, torch_dtype=DTYPE)
+    pipe.enable_model_cpu_offload(device=DEVICE)
+
+    # Stage 1 — cheap draft at half resolution.
+    print(f"1/3  Stage 1 @ {half_w}×{half_h}")
+    video_latent, audio_latent = pipe(
+        prompt=prompt,
+        negative_prompt=DEFAULT_NEGATIVE_PROMPT,
+        height=half_h,
+        width=half_w,
+        num_frames=NUM_FRAMES,
+        frame_rate=FPS,
+        num_inference_steps=len(DISTILLED_SIGMA_VALUES),
+        sigmas=DISTILLED_SIGMA_VALUES,
+        guidance_scale=1.0,
+        generator=generator,
+        output_type="latent",
+        return_dict=False,
+    )
+
+    # Upsample — bring the latent to full resolution before refining.
+    print(f"2/3  Upsample → {WIDTH}×{HEIGHT}")
+    upsampler = LTX2LatentUpsamplerModel.from_pretrained(
+        SPATIAL_UPSAMPLER, subfolder="latent_upsampler", torch_dtype=DTYPE
+    )
+    upsample_pipe = LTX2LatentUpsamplePipeline(vae=pipe.vae, latent_upsampler=upsampler)
+    upsample_pipe.enable_model_cpu_offload(device=DEVICE)
+    video_latent = upsample_pipe(
+        latents=video_latent[:1],
+        height=half_h,
+        width=half_w,
+        num_frames=NUM_FRAMES,
+        output_type="latent",
+        return_dict=False,
+    )[0]
+    del upsample_pipe, upsampler
+    torch.cuda.empty_cache()
+
+    # Stage 2 — refine at full resolution (renoises from STAGE_2 schedule).
+    print(f"3/3  Stage 2 @ {WIDTH}×{HEIGHT}")
+    video_latent, _ = pipe(
+        latents=video_latent,
+        audio_latents=audio_latent,
+        prompt=prompt,
+        negative_prompt=DEFAULT_NEGATIVE_PROMPT,
+        height=HEIGHT,
+        width=WIDTH,
+        num_frames=NUM_FRAMES,
+        frame_rate=FPS,
+        num_inference_steps=len(STAGE_2_DISTILLED_SIGMA_VALUES),
+        noise_scale=STAGE_2_DISTILLED_SIGMA_VALUES[0],
+        sigmas=STAGE_2_DISTILLED_SIGMA_VALUES,
+        guidance_scale=1.0,
+        generator=generator,
+        output_type="latent",
+        return_dict=False,
+    )
+
+    latent = video_latent.detach().cpu()
+    del pipe
+    torch.cuda.empty_cache()
+    return latent
+
+
+def compare_decoders(latent: torch.Tensor, out_dir: Path) -> None:
+    """Decode the same latent with LTX-2.3 and PrunaVAED; print ms and save mp4s."""
+    # Needed so Diffusers can build the pruned decoder graph correctly.
+    patch_pruna_ltx2_decoder()
+
+    results = {}
+    for name, model_id in (("ltx23", LTX23_VAE), ("prunavaed", str(PRUNED_VAE))):
+        print(f"Decoding with {name} …")
+        vae = load_vae(model_id)
+        video, ms = timed_decode(vae, latent)
+        results[name] = ms
+        path = out_dir / f"{name}.mp4"
+        save_mp4(video, path)
+        print(f"  {name}: {ms:.1f} ms  →  {path}")
+        del vae, video
+        torch.cuda.empty_cache()
+
+    print(f"Speedup: {results['ltx23'] / results['prunavaed']:.2f}×")
+
+
+# ---------------------------------------------------------------------------
+# Entry point
+# ---------------------------------------------------------------------------
+
+def main() -> None:
+    if not torch.cuda.is_available():
+        raise SystemExit("This demo needs a CUDA GPU.")
+
+    print(f"Prompt: {PROMPT[:80]}…")
+    print(f"Output: {OUTPUT_DIR}")
+
+    latent = generate_latent(PROMPT)
+    print(f"Latent shape: {tuple(latent.shape)}")
+    compare_decoders(latent, OUTPUT_DIR)
+
+
+if __name__ == "__main__":
+    main()