Refining demo

demo/demo_distilled_decode.py

@@ -2,7 +2,8 @@
 """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.
+(PrunaVAED) on the *same* latent. Prints decode time (ms) and peak VRAM (GiB),
+and writes two mp4s.
 
 Requires a CUDA GPU. From the repo root:
 
@@ -28,24 +29,25 @@ from diffusers.pipelines.ltx2.utils import (
     STAGE_2_DISTILLED_SIGMA_VALUES,
 )
 
+REPO_ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(REPO_ROOT))
+
+from patch_diffusers import patch_pruna_ltx2_decoder  # noqa: E402
+
+
 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
+PRUNED_VAE = str(REPO_ROOT)  # this repo's vae/ folder
 
-# ~720p for 5 s @ 24 fps. Height/width must be multiples of 64 (2-stage).
+# ~1080p 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
@@ -63,46 +65,51 @@ OUTPUT_DIR = REPO_ROOT / "outputs" / "demo_distilled"
 # Helpers
 # ---------------------------------------------------------------------------
 
-def save_mp4(video: torch.Tensor, path: Path) -> None:
+def save_mp4(video: torch.Tensor, path: Path, fps: float) -> 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()
+    frames = video[0].permute(1, 2, 3, 0).clamp(-1, 1).float()
+    frames = ((frames + 1) / 2 * 255).round().byte().cpu().numpy()
     path.parent.mkdir(parents=True, exist_ok=True)
-    iio.imwrite(path, frames, fps=FPS, codec="libx264")
+    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."""
+    """Load a video VAE decoder on GPU, tiling off so we time full-frame decodes."""
     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
+    vae.disable_tiling()
+    return vae.to(DEVICE).eval()
 
 
 @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)."""
+def timed_decode(
+    vae: AutoencoderKLLTX2Video, latent: torch.Tensor
+) -> tuple[torch.Tensor, float, float]:
+    """Decode DECODE_RUNS times after warm-up.
+
+    Returns the video on CPU, the median latency in ms and the peak VRAM in GiB.
+    Deliberately no autocast: weights and latent are already bfloat16, and autocast
+    would run every PerChannelRMSNorm (``x**2``) in float32, costing time and VRAM.
+    """
     latent = latent.to(DEVICE, DTYPE)
 
-    with torch.autocast(DEVICE, dtype=DTYPE):
-        for _ in range(DECODE_WARMUP):
-            vae.decode(latent, return_dict=False)
+    for _ in range(DECODE_WARMUP):
+        vae.decode(latent, return_dict=False)
 
+    torch.cuda.synchronize()
+    torch.cuda.reset_peak_memory_stats()
     times_ms = []
-    video = None
     for _ in range(DECODE_RUNS):
+        video = None  # keep the previous output out of the peak VRAM measurement
         torch.cuda.synchronize()
         t0 = time.perf_counter()
-        with torch.autocast(DEVICE, dtype=DTYPE):
-            video = vae.decode(latent, return_dict=False)[0]
+        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)
+    peak_gib = torch.cuda.max_memory_allocated() / 2**30
+    return video.cpu(), statistics.median(times_ms), peak_gib
 
 
 @torch.inference_mode()
@@ -180,23 +187,24 @@ def generate_latent(prompt: str) -> torch.Tensor:
 
 
 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."""
+    """Decode the same latent with LTX-2.3 and PrunaVAED; print ms/VRAM 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))):
+    for name, model_id in (("ltx23", LTX23_VAE), ("prunavaed", PRUNED_VAE)):
         print(f"Decoding with {name} …")
         vae = load_vae(model_id)
-        video, ms = timed_decode(vae, latent)
-        results[name] = ms
+        video, ms, peak_gib = timed_decode(vae, latent)
+        results[name] = (ms, peak_gib)
         path = out_dir / f"{name}.mp4"
-        save_mp4(video, path)
-        print(f"  {name}: {ms:.1f} ms  →  {path}")
+        save_mp4(video, path, FPS)
+        print(f"  {name}: {ms:.1f} ms · {peak_gib:.2f} GiB peak  →  {path}")
         del vae, video
         torch.cuda.empty_cache()
 
-    print(f"Speedup: {results['ltx23'] / results['prunavaed']:.2f}×")
+    (ltx_ms, ltx_gib), (pruna_ms, pruna_gib) = results["ltx23"], results["prunavaed"]
+    print(f"Speedup: {ltx_ms / pruna_ms:.2f}×    Peak VRAM: {pruna_gib / ltx_gib:.0%} of LTX-2.3")
 
 
 # ---------------------------------------------------------------------------