demo/demo_distilled_decode.py

#!/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 peak VRAM (GiB),
and writes two mp4s.

Requires a CUDA GPU. From the repo root:

    pip install -r requirements-test.txt
    python demo/demo_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,
)

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)
# ---------------------------------------------------------------------------

# 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 = str(REPO_ROOT)  # this repo's vae/ folder

# ~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

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" / "demo_distilled"


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

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).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")


def load_vae(model_id: str) -> AutoencoderKLLTX2Video:
    """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.disable_tiling()
    return vae.to(DEVICE).eval()


@torch.inference_mode()
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)

    for _ in range(DECODE_WARMUP):
        vae.decode(latent, return_dict=False)

    torch.cuda.synchronize()
    torch.cuda.reset_peak_memory_stats()
    times_ms = []
    for _ in range(DECODE_RUNS):
        video = None  # keep the previous output out of the peak VRAM measurement
        torch.cuda.synchronize()
        t0 = time.perf_counter()
        video = vae.decode(latent, return_dict=False)[0]
        torch.cuda.synchronize()
        times_ms.append((time.perf_counter() - t0) * 1000)

    peak_gib = torch.cuda.max_memory_allocated() / 2**30
    return video.cpu(), statistics.median(times_ms), peak_gib


@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/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", PRUNED_VAE)):
        print(f"Decoding with {name} …")
        vae = load_vae(model_id)
        video, ms, peak_gib = timed_decode(vae, latent)
        results[name] = (ms, peak_gib)
        path = out_dir / f"{name}.mp4"
        save_mp4(video, path, FPS)
        print(f"  {name}: {ms:.1f} ms · {peak_gib:.2f} GiB peak  →  {path}")
        del vae, video
        torch.cuda.empty_cache()

    (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")


# ---------------------------------------------------------------------------
# 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()