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)