app.py

import os

# Must be set before torch initialises its caching allocator: the causal VAE
# decode of 81 frames at 480x832 fragments the pool badly otherwise.
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import spaces  # noqa: F401  -- MUST come before any torch / CUDA-touching import

import random
import tempfile
import time
import uuid
from pathlib import Path

import numpy as np
import torch
import gradio as gr
from omegaconf import OmegaConf
from huggingface_hub import snapshot_download

import av

from demo_utils.constant import ZERO_VAE_CACHE
from demo_utils.vae_block3 import VAEDecoderWrapper
from pipeline import CausalInferencePipeline
from utils.wan_wrapper import WanDiffusionWrapper, WanTextEncoder

# ---------------------------------------------------------------------------
# Checkpoints
# ---------------------------------------------------------------------------
BASE_MODEL_ID = "Wan-AI/Wan2.1-T2V-1.3B"
STUDENT_MODEL_ID = "LiJiaxing/DistillAlign"
STUDENT_CKPT = "checkpoints/distillalign_distill_1p3b_teacher.pt"
MODEL_ROOT = "wan_models"
MODEL_NAME = "Wan2.1-T2V-1.3B"
MODEL_DIR = f"{MODEL_ROOT}/{MODEL_NAME}"

torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
torch.set_grad_enabled(False)

print("== downloading Wan2.1-T2V-1.3B VAE + UMT5-XXL text encoder ==", flush=True)
snapshot_download(
    repo_id=BASE_MODEL_ID,
    local_dir=MODEL_DIR,
    repo_type="model",
    allow_patterns=[
        "Wan2.1_VAE.pth",
        "models_t5_umt5-xxl-enc-bf16.pth",
        # CausalWanModel.from_pretrained() builds the backbone from this file's
        # weights before the DistillAlign student state dict replaces them, so
        # the base transformer has to be on disk even though it is discarded.
        "diffusion_pytorch_model.safetensors",
        "google/umt5-xxl/*",
        "config.json",
    ],
)

print("== downloading the DistillAlign 1.3B distilled student ==", flush=True)
snapshot_download(
    repo_id=STUDENT_MODEL_ID,
    local_dir=".",
    repo_type="model",
    allow_patterns=[STUDENT_CKPT, "config.json"],
)

# ---------------------------------------------------------------------------
# Config + models (module scope, eagerly on the GPU -- ZeroGPU rule #2)
# ---------------------------------------------------------------------------
config = OmegaConf.merge(
    OmegaConf.load("configs/default_config.yaml"),
    OmegaConf.load("configs/self_forcing_dmd.yaml"),
)
config.model_kwargs.model_root = MODEL_ROOT
config.model_kwargs.model_name = MODEL_NAME


def _load_checkpoint(path: Path, state_key: str = "auto"):
    """Mirror of ``distribution_eval.sampling._load_checkpoint`` (Apache-2.0)."""
    try:
        obj = torch.load(path, mmap=True, map_location="cpu", weights_only=True)
    except TypeError:
        obj = torch.load(path, map_location="cpu", weights_only=True)
    if not isinstance(obj, dict):
        raise TypeError(f"checkpoint must contain a state dict, got {type(obj)}")

    selected = state_key
    if state_key == "auto":
        if "generator_ema" in obj:
            selected = "generator_ema"
        elif "generator" in obj:
            selected = "generator"
        else:
            selected = "root"

    state = obj if selected == "root" else obj.get(selected)
    if not isinstance(state, dict):
        raise KeyError(f"checkpoint has no state dict named {selected!r}")

    cleaned = {}
    for key, value in state.items():
        if key.startswith("model._fsdp_wrapped_module."):
            key = key.replace("model._fsdp_wrapped_module.", "model.", 1)
        elif key.startswith("_fsdp_wrapped_module."):
            key = key.replace("_fsdp_wrapped_module.", "", 1)
        cleaned[key] = value
    return cleaned, selected


print("== loading UMT5-XXL text encoder ==", flush=True)
text_encoder = WanTextEncoder(model_name=MODEL_NAME, model_root=MODEL_ROOT)
text_encoder.eval().requires_grad_(False)
text_encoder.to(device="cuda", dtype=torch.bfloat16)

print("== loading the causal Wan2.1-1.3B student ==", flush=True)
transformer = WanDiffusionWrapper(
    model_name=MODEL_NAME,
    model_root=MODEL_ROOT,
    timestep_shift=float(config.timestep_shift),
    is_causal=True,
).eval()
_state_dict, _state_key = _load_checkpoint(Path(STUDENT_CKPT), "auto")
print(f"== loading DistillAlign state dict '{_state_key}' ==", flush=True)
transformer.load_state_dict(_state_dict, strict=True)
del _state_dict
transformer.requires_grad_(False)
transformer.to(dtype=torch.bfloat16).to("cuda")

print("== loading streaming (causal) VAE decoder ==", flush=True)
vae = VAEDecoderWrapper().eval()
_vae_sd = torch.load(f"{MODEL_DIR}/Wan2.1_VAE.pth", map_location="cpu", weights_only=True)
vae.load_state_dict({k: v for k, v in _vae_sd.items() if "decoder." in k or "conv2" in k})
del _vae_sd
vae.requires_grad_(False)
vae.to(dtype=torch.bfloat16).to("cuda")

pipeline = CausalInferencePipeline(
    config,
    device="cuda",
    generator=transformer,
    text_encoder=text_encoder,
    vae=vae,
).to(dtype=torch.bfloat16)

LATENT_SHAPE = (1, 21, 16, 60, 104)  # 21 latent frames -> 81 pixel frames @ 480x832
NUM_BLOCKS = LATENT_SHAPE[1] // pipeline.num_frame_per_block
FPS = 16
torch.cuda.empty_cache()
print(
    f"== ready: {NUM_BLOCKS} blocks x {pipeline.num_frame_per_block} latent frames, "
    f"denoising steps {pipeline.denoising_step_list.tolist()}, "
    f"{torch.cuda.memory_allocated() / 1024 ** 3:.1f} GB resident ==",
    flush=True,
)


# ---------------------------------------------------------------------------
# Video encoding
# ---------------------------------------------------------------------------
def frames_to_mp4(frames, filepath, fps=FPS):
    """frames: list of HWC uint8 RGB arrays -> h264 mp4."""
    height, width = frames[0].shape[:2]
    container = av.open(filepath, mode="w", format="mp4")
    stream = container.add_stream("h264", rate=fps)
    stream.width = width
    stream.height = height
    stream.pix_fmt = "yuv420p"
    stream.options = {"crf": "18", "preset": "veryfast"}
    try:
        for frame_np in frames:
            frame = av.VideoFrame.from_ndarray(np.ascontiguousarray(frame_np), format="rgb24")
            for packet in stream.encode(frame):
                container.mux(packet)
        for packet in stream.encode():
            container.mux(packet)
    finally:
        container.close()
    return filepath


# ---------------------------------------------------------------------------
# Inference
# ---------------------------------------------------------------------------
@torch.no_grad()
# Measured on a ZeroGPU (zero-a10g / RTX PRO 6000) worker: ~15.5 s of GPU time
# end-to-end (7 causal blocks x 4 DMD steps + streaming causal VAE decode + mp4
# mux), so 20 s is the observed time plus the usual ~15% margin. Keeping this
# tight matters: ZeroGPU compares the requested duration against the visitor's
# remaining quota and a large request also lowers the demo's queue priority.
@spaces.GPU(duration=20)
def generate_video(prompt: str, seed: int = 0) -> tuple[str, str]:
    """Generate a 5-second 480x832 video from a text prompt.

    Args:
        prompt: Text description of the video to generate.
        seed: Random seed. Use 0 for a random seed.

    Returns:
        The path to the generated mp4, plus a short status line.
    """
    if not prompt or not prompt.strip():
        raise gr.Error("Please enter a text prompt.")

    prompt = prompt.strip()
    if seed is None:
        seed = 0
    seed = int(seed)
    if seed == 0:
        seed = random.randint(1, 2**31 - 1)

    t0 = time.perf_counter()
    device = "cuda"

    conditional_dict = text_encoder(text_prompts=[prompt])
    for key, value in conditional_dict.items():
        conditional_dict[key] = value.to(dtype=torch.bfloat16)

    rnd = torch.Generator(device=device).manual_seed(seed)
    noise = torch.randn(LATENT_SHAPE, device=device, dtype=torch.bfloat16, generator=rnd)

    pipeline._initialize_kv_cache(1, torch.bfloat16, device=device)
    pipeline._initialize_crossattn_cache(1, torch.bfloat16, device=device)
    vae_cache = [c.to(device=device, dtype=torch.bfloat16) for c in ZERO_VAE_CACHE]

    denoising_step_list = pipeline.denoising_step_list
    all_frames = []
    current_start_frame = 0

    for block_index in range(NUM_BLOCKS):
        current_num_frames = pipeline.num_frame_per_block
        noisy_input = noise[:, current_start_frame:current_start_frame + current_num_frames]

        for step_index, current_timestep in enumerate(denoising_step_list):
            timestep = torch.ones(
                [1, current_num_frames], device=device, dtype=torch.int64
            ) * current_timestep
            _, denoised_pred = pipeline.generator(
                noisy_image_or_video=noisy_input,
                conditional_dict=conditional_dict,
                timestep=timestep,
                kv_cache=pipeline.kv_cache1,
                crossattn_cache=pipeline.crossattn_cache,
                current_start=current_start_frame * pipeline.frame_seq_length,
            )
            if step_index < len(denoising_step_list) - 1:
                next_timestep = denoising_step_list[step_index + 1]
                noisy_input = pipeline.scheduler.add_noise(
                    denoised_pred.flatten(0, 1),
                    torch.randn_like(denoised_pred.flatten(0, 1)),
                    next_timestep * torch.ones(
                        [1 * current_num_frames], device=device, dtype=torch.long
                    ),
                ).unflatten(0, denoised_pred.shape[:2])

        # Re-run with clean context so the KV cache holds denoised latents for
        # the next block (skipped on the last block -- the cache is discarded).
        if block_index < NUM_BLOCKS - 1:
            pipeline.generator(
                noisy_image_or_video=denoised_pred,
                conditional_dict=conditional_dict,
                timestep=torch.zeros_like(timestep),
                kv_cache=pipeline.kv_cache1,
                crossattn_cache=pipeline.crossattn_cache,
                current_start=current_start_frame * pipeline.frame_seq_length,
            )

        # Stream the VAE decode one block at a time to bound peak memory. The
        # causal decoder emits 4 pixel frames per latent frame; the very first
        # block carries 3 warm-up frames from the causal padding, so drop them
        # (9 + 6*12 = 81 pixel frames, matching the reference decode()).
        pixels, vae_cache = pipeline.vae(denoised_pred.to(torch.bfloat16), *vae_cache)
        if block_index == 0:
            pixels = pixels[:, 3:]
        pixels = pixels.float()
        for frame in pixels[0]:
            arr = torch.clamp(frame, -1.0, 1.0).add(1.0).mul(127.5)
            all_frames.append(arr.to(torch.uint8).permute(1, 2, 0).cpu().numpy())

        current_start_frame += current_num_frames
        print(
            f"block {block_index + 1}/{NUM_BLOCKS} decoded, "
            f"{len(all_frames)} frames, {time.perf_counter() - t0:.1f}s",
            flush=True,
        )

    out_path = os.path.join(tempfile.gettempdir(), f"distillalign-{uuid.uuid4().hex}.mp4")
    frames_to_mp4(all_frames, out_path, fps=FPS)

    elapsed = time.perf_counter() - t0
    status = (
        f"Generated {len(all_frames)} frames at {FPS} fps "
        f"({len(all_frames) / FPS:.1f}s of video) in {elapsed:.1f}s. Seed: {seed}."
    )
    print(status, flush=True)
    return out_path, status


# ---------------------------------------------------------------------------
# UI
# ---------------------------------------------------------------------------
EXAMPLE_PROMPTS = [
    "Four teenage girls, with windswept hair and determined expressions, ride powerful waves in "
    "frigid, crystal-clear coldwater. Their vibrant surfboards cut through frothy whitecaps as they "
    "lean into turns with athletic grace. Dynamic camera angles—low drone sweeps, close-up shoulder "
    "shots, and overhead tracking—capture their synchronized movements. Each girl wears snug "
    "wetsuits, laughing and shouting as they paddle, carve, and ride the swell. Motion is fluid and "
    "energetic, emphasizing their youthful spirit against the vast, chilly ocean backdrop.",
    "A young couple stands side by side under a vast, starry night sky, gazing upward in wonder as "
    "distant stars flicker with soft, ethereal light. Their faces glow with awe, hands lightly "
    "touching, shoulders relaxed. The camera slowly pans around them, tilting upward to capture the "
    "infinite cosmos, then gently zooms in on their eyes reflecting the stars. Shot in cinematic 4K "
    "HD, with deep blues and warm golden hues, the scene pulses with quiet romance and cosmic wonder "
    "as gentle wind rustles their hair.",
    "A cinematic, slow-motion sequence capturing every chess move in vivid detail, filmed in 16:9 "
    "aspect ratio with dynamic camera tracking following each piece's movement across the board. The "
    "chess pieces are intricately designed, with dramatic lighting casting sharp shadows. The camera "
    "pans smoothly between moves, zooms in for critical captures, and tilts slightly to emphasize "
    "strategic tension. Each piece glides with purpose — pawns advance, knights leap, bishops slash "
    "diagonally — all rendered with fluid motion and cinematic grace.",
    "Hyper realistic cinematic shot of vibrant orange fish gracefully leaping and floating on "
    "shimmering water surface, capturing dynamic ripples and vivid reflections. Fish vary in size, "
    "displaying lively expressions as they dart through the frame. Camera glides smoothly at low "
    "angle, panning across the water's surface to emphasize motion and depth. Rich, saturated colors "
    "contrast with soft ambient lighting. Immersive and fluid, evoking a dreamlike aquatic ballet "
    "with natural, organic movement.",
    "A man with long, flowing curly black hair stands alone on a windswept beach cliff, gazing at "
    "the sun setting over the endless ocean — captured from behind in a stunning anime-style digital "
    "illustration by Makoto Shinkai. His hollow cheeks and contemplative posture evoke melancholy "
    "beauty, as if frozen in his last moment of solitude. The scene glows with warm, cinematic hues, "
    "enhanced by soft atmospheric lighting and gentle waves crashing below. Camera slowly pans right, "
    "then tilts up to frame the sky, blending romanticism with aestheticism. Motion: slow, "
    "reflective, emotionally resonant.",
    "A dynamic slow-motion shot of a figure skater gliding across an ice rink, arms elegantly "
    "extended, showcasing fluid spins and graceful landings. His athletic build, focused expression, "
    "and crisp costume details are highlighted under bright arena lights. Camera circles him "
    "mid-spin, then tracks his glide with a low-angle follow shot. Emphasize the shimmering ice, "
    "crisp blades, and swirling motion trails. Cinematic, high-definition, with dramatic lighting and "
    "smooth camera transitions capturing the intensity and artistry of competitive figure skating.",
]

CSS = """
#col-container { max-width: 1100px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

with gr.Blocks(title="DistillAlign", theme=gr.themes.Citrus(), css=CSS) as demo:
    gr.Markdown(
        """
        # 🎬 DistillAlign — Autoregressive Video Distillation
        **Coordinating Mode Covering and Mode Seeking in Autoregressive Video Distillation**

        A **Wan2.1-T2V-1.3B** student distilled for few-step *causal* (autoregressive) text-to-video
        generation. It samples 21 latent frames in 7 causal blocks with a 4-step DMD schedule
        (`1000 → 750 → 500 → 250`) and a KV cache, then streams them through the causal Wan VAE —
        **81 frames / 480×832 / 16 fps ≈ 5 s** in a single forward pass per block.

        ⏱️ Generation takes roughly a minute on an RTX PRO 6000 (ZeroGPU). Pick one of the example
        prompts below or write your own.
        """
    )

    with gr.Row():
        with gr.Column(scale=1):
            prompt = gr.Textbox(
                label="Prompt",
                placeholder="Describe the video you want to generate...",
                lines=6,
                max_lines=10,
            )
            generate_btn = gr.Button("Generate video", variant="primary", size="lg")
            with gr.Accordion("Advanced options", open=False):
                seed = gr.Number(
                    label="Seed (0 = random)", value=0, precision=0, minimum=0
                )
        with gr.Column(scale=1):
            video = gr.Video(label="Generated video", autoplay=True, height=420)
            status = gr.Markdown()

    generate_btn.click(
        fn=generate_video,
        inputs=[prompt, seed],
        outputs=[video, status],
        api_name="generate",
    )

    gr.Examples(
        examples=[[p] for p in EXAMPLE_PROMPTS],
        inputs=[prompt],
        outputs=[video, status],
        fn=generate_video,
        cache_examples=True,
        cache_mode="lazy",
        label="Example prompts from the DistillAlign paper",
    )

    gr.Markdown(
        """
        ---
        **Model**: [LiJiaxing/DistillAlign](https://huggingface.co/LiJiaxing/DistillAlign) ·
        **Code**: [LiJiaxing0213/DistillAlign](https://github.com/LiJiaxing0213/DistillAlign) ·
        **Paper**: [arXiv 2607.26811](https://huggingface.co/papers/2607.26811)

        Built on Wan2.1-T2V-1.3B ([Wan-AI](https://huggingface.co/Wan-AI/Wan2.1-T2V-1.3B)).
        The example prompts above are the authors' own evaluation prompts
        (`prompts/distribution_eval_16.txt`, Apache-2.0).
        """
    )

if __name__ == "__main__":
    demo.queue().launch(mcp_server=True)