"""AAD-1: Asymmetric Adversarial Distillation for One-Step Autoregressive Video Generation. A Gradio demo that loads the AAD-1 one-step image-to-video model and generates short video clips from a reference image and a text prompt. """ import os os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") os.environ.setdefault("TOKENIZERS_PARALLELISM", "false") import spaces # MUST come before torch / any CUDA-touching import import gc import json import tempfile import time from pathlib import Path from types import SimpleNamespace import gradio as gr import numpy as np import torch from PIL import Image # ---- repo-local imports (the AAD-1 source tree is uploaded alongside app.py) ---- from aad1.checkpoint_io import ( load_native_generator_checkpoint_into_model, ) from pipeline import CausalInferencePipeline from utils.misc import set_seed from utils.wan_wrapper import WanDiffusionWrapper, WanTextEncoder, WanVAEWrapper from huggingface_hub import snapshot_download # --------------------------------------------------------------------------- # Constants # --------------------------------------------------------------------------- MODEL_ID_AAD1 = "Watay/AAD-1" MODEL_ID_WAN = "Wan-AI/Wan2.1-T2V-14B" IMAGE_HEIGHT = 480 IMAGE_WIDTH = 832 DEFAULT_NUM_FRAMES = 81 # 5s @ 16fps → 81 frames (16n+1 rule) DEFAULT_FRAME_RATE = 16 DEFAULT_SEED = 1000 DEFAULT_LOCAL_ATTN_SIZE = 9 DEFAULT_SINK_SIZE = 1 # --------------------------------------------------------------------------- # Model loading (module scope, eager .to("cuda") for ZeroGPU pack) # --------------------------------------------------------------------------- def _download_weights(): """Download AAD-1 checkpoint shards and the shared Wan2.1-T2V-14B model dir.""" from huggingface_hub import hf_hub_download print("[aad-1] Downloading AAD-1 checkpoint …") aad1_dir = snapshot_download( repo_id=MODEL_ID_AAD1, allow_patterns=["14b_i2v_1step_transformer/*"], repo_type="model", ) # Use the snapshot path directly — do NOT resolve symlinks, because # checkpoint_io's ensure_native_checkpoint_path resolves and checks the # suffix. The blob path has no .index.json suffix. checkpoint_path = Path(aad1_dir) / "14b_i2v_1step_transformer" / "self_forcing_generator_bf16.index.json" print("[aad-1] Downloading shared Wan2.1-T2V-14B model dir …") wan_model_dir = snapshot_download( repo_id=MODEL_ID_WAN, allow_patterns=[ "config.json", "Wan2.1_VAE.pth", "models_t5_umt5-xxl-enc-bf16.pth", "google/umt5-xxl/*", ], repo_type="model", ) return checkpoint_path, Path(wan_model_dir) def _build_runtime_config(num_frames=DEFAULT_NUM_FRAMES): """Build the SimpleNamespace config consumed by the pipeline.""" if (num_frames - 1) % 16 != 0: raise ValueError(f"`num_frames` must satisfy 16n+1, got {num_frames}.") latent_frames = (num_frames - 1) // 4 + 1 return SimpleNamespace( generator_name="Wan2.1-T2V-14B", text_len=512, model_kwargs={ "timestep_shift": 5.0, "local_attn_size": DEFAULT_LOCAL_ATTN_SIZE, "sink_size": DEFAULT_SINK_SIZE, }, denoising_step_list=[1000], warp_denoising_step=True, image_or_video_shape=[1, latent_frames, 16, IMAGE_HEIGHT // 8, IMAGE_WIDTH // 8], num_training_frames=latent_frames, num_frame_per_block=4, independent_first_frame=True, mixed_precision=True, context_noise=0, ) print("[aad-1] Starting model load …") _t0 = time.perf_counter() checkpoint_path, wan_model_dir = _download_weights() # Don't call ensure_native_checkpoint_path — it resolves symlinks to blob # paths which lose the .index.json suffix. Just verify the file exists. if not checkpoint_path.exists(): raise FileNotFoundError(f"Checkpoint index not found: {checkpoint_path}") _config = _build_runtime_config(DEFAULT_NUM_FRAMES) # Text encoder (T5 umt5-xxl) — kept on CPU to save VRAM, moved to GPU per-call text_encoder = WanTextEncoder( model_name=_config.generator_name, model_dir=wan_model_dir, compute_dtype=torch.float32, output_dtype=torch.bfloat16, ) text_encoder.eval() text_encoder.requires_grad_(False) # VAE vae = WanVAEWrapper(_config.generator_name, model_dir=wan_model_dir) # Diffusion transformer (causal) transformer = WanDiffusionWrapper( model_name=_config.generator_name, model_config=_config, is_causal=True, model_dir=wan_model_dir, **_config.model_kwargs, ) load_native_generator_checkpoint_into_model(transformer.model, checkpoint_path, strict=True) transformer.eval() transformer.requires_grad_(False) # Pipeline pipeline = CausalInferencePipeline( _config, device=torch.device("cuda"), generator=transformer, text_encoder=text_encoder, vae=vae, ) # Move transformer to "cuda" (ZeroGPU intercepts this for packing) transformer.to("cuda") print(f"[aad-1] Model load completed in {time.perf_counter() - _t0:.1f}s") # --------------------------------------------------------------------------- # Inference # --------------------------------------------------------------------------- def _load_image_tensor(image: Image.Image) -> torch.Tensor: """Resize and normalise a PIL image to the model's expected tensor format.""" image = image.convert("RGB").resize((IMAGE_WIDTH, IMAGE_HEIGHT), Image.LANCZOS) arr = np.asarray(image, dtype=np.float32) / 255.0 arr = arr * 2.0 - 1.0 # normalise to [-1, 1] tensor = torch.from_numpy(arr).permute(2, 0, 1).unsqueeze(0).unsqueeze(2) return tensor.contiguous() def _save_video(video: torch.Tensor, fps: int = DEFAULT_FRAME_RATE) -> str: """Save a [B, T, C, H, W] tensor (values in [0, 1]) to a temporary mp4 file.""" import imageio.v2 as imageio frames = (255.0 * video).clamp(0, 255).to(torch.uint8) frames = frames.permute(0, 1, 3, 4, 2).cpu().numpy() # [B, T, H, W, C] tmp = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) tmp.close() with imageio.get_writer(tmp.name, fps=fps, codec="libx264", format="FFMPEG") as writer: for frame in frames[0]: writer.append_data(frame) return tmp.name @spaces.GPU(duration=180) def generate_video( image: Image.Image, prompt: str, num_frames: int = DEFAULT_NUM_FRAMES, seed: int = DEFAULT_SEED, fps: int = DEFAULT_FRAME_RATE, ) -> str: """Generate a short video from a reference image and a text prompt. Args: image: The reference / conditioning image. prompt: A text prompt describing the desired motion and scene. num_frames: Number of output frames (must satisfy 16n+1; default 81 ≈ 5s). seed: Random seed for reproducibility. fps: Output video frame rate. """ if image is None: raise gr.Error("Please provide a reference image.") if not prompt or not prompt.strip(): raise gr.Error("Please provide a text prompt.") # Validate num_frames valid_frames = [17, 33, 49, 65, 81, 97, 113, 129, 145, 161] if num_frames not in valid_frames: # Snap to nearest valid value num_frames = min(valid_frames, key=lambda x: abs(x - num_frames)) config = _build_runtime_config(num_frames) pipeline.args = config pipeline.generator.model_config = config pipeline.denoising_step_list = torch.tensor(config.denoising_step_list, dtype=torch.long) if config.warp_denoising_step: timesteps = torch.cat((pipeline.scheduler.timesteps.cpu(), torch.tensor([0], dtype=torch.float32))) pipeline.denoising_step_list = timesteps[1000 - pipeline.denoising_step_list] pipeline.gradient_frames = config.image_or_video_shape[1] pipeline.generator.seq_len = ( config.image_or_video_shape[1] * config.image_or_video_shape[3] * config.image_or_video_shape[4] // int(np.prod(pipeline.patch_size)) ) pipeline.generator.max_attention_size = pipeline.generator.seq_len pipeline.token_seq_length = ( config.image_or_video_shape[1] * config.image_or_video_shape[3] * config.image_or_video_shape[4] // int(np.prod(pipeline.patch_size)) ) if pipeline.local_attn_size != -1: pipeline.kv_cache_size = pipeline.local_attn_size * pipeline.frame_seq_length else: pipeline.kv_cache_size = pipeline.token_seq_length gen_model = getattr(pipeline.generator.model, "_orig_mod", pipeline.generator.model) gen_model.num_frame_per_block = config.num_frame_per_block gen_model.independent_first_frame = config.independent_first_frame gen_model.max_attention_size = pipeline.generator.max_attention_size # Reset block mask when frame count changes cached_frames = getattr(gen_model, "_aad1_demo_block_mask_frames", None) if cached_frames != config.image_or_video_shape[1]: gen_model.block_mask = None gen_model._aad1_demo_block_mask_frames = config.image_or_video_shape[1] set_seed(seed) torch.set_grad_enabled(False) # --- Text encoding --- conditional_dict = text_encoder([prompt]) conditional_dict = { k: v.to(device="cuda") if torch.is_tensor(v) else v for k, v in conditional_dict.items() } # --- Image latent --- image_tensor = _load_image_tensor(image).to(device="cuda", dtype=torch.float32) pipeline.vae.to("cuda") initial_latent = pipeline.vae.encode_to_latent(image_tensor).to(device="cuda") torch.cuda.empty_cache() # --- Noise --- noise_generator = torch.Generator("cpu").manual_seed(seed) latent_frames = config.image_or_video_shape[1] latent_channels = config.image_or_video_shape[2] latent_h = config.image_or_video_shape[3] latent_w = config.image_or_video_shape[4] full_noise_bcthw = torch.randn( [1, latent_channels, latent_frames, latent_h, latent_w], generator=noise_generator, dtype=torch.float32, ) full_noise = full_noise_bcthw.permute(0, 2, 1, 3, 4).contiguous() sampled_noise = full_noise[:, initial_latent.shape[1]:].to(device="cuda") # --- Run inference --- video, _ = pipeline.inference( noise=sampled_noise, text_prompts=None, return_latents=True, initial_latent=initial_latent, noise_generator=noise_generator, low_memory=True, conditional_dict=conditional_dict, ) # --- Clean up --- pipeline.kv_cache1 = None pipeline.crossattn_cache = None pipeline.vae.to("cpu") del conditional_dict gc.collect() torch.cuda.empty_cache() # --- Save --- video_path = _save_video(video.cpu(), fps=fps) return video_path # --------------------------------------------------------------------------- # Gradio UI # --------------------------------------------------------------------------- CSS = """ #col-container { max-width: 1100px; margin: 0 auto; } .dark .gradio-container { color: var(--body-text-color); } """ EXAMPLES_DIR = Path(__file__).resolve().parent / "examples" EXAMPLES = [ [str(EXAMPLES_DIR / "horses_running_dirt.jpg"), "a couple of horses are running in the dirt", 81, 1000, 16], [str(EXAMPLES_DIR / "cat_dog.png"), "A lively scene featuring a brown and white dog energetically chasing after a sleek black cat through a grassy backyard. The dog is running with its tongue out, tail wagging, and ears perked up, while the cat is sprinting swiftly, occasionally glancing over its shoulder with alert eyes.", 81, 42, 16], [str(EXAMPLES_DIR / "beach.png"), "A serene sunset over the pristine beaches of Cancun, Mexico. The sky is painted with vibrant hues of orange, pink, and purple, reflecting off the calm turquoise waters.", 81, 100, 16], [str(EXAMPLES_DIR / "blue_smoke.jpg"), "a blue and white smoke is swirly in the dark", 81, 999, 16], [str(EXAMPLES_DIR / "viking.png"), "A Viking warrior in full armor and a horned helmet, sitting at a wooden table, is enthusiastically eating a large, modern-style pizza topped with pepperoni and cheese.", 81, 7, 16], ] with gr.Blocks(theme=gr.themes.Citrus(), css=CSS) as demo: with gr.Column(elem_id="col-container"): gr.Markdown("# AAD-1: One-Step Autoregressive Video Generation") gr.Markdown( "Generate short videos from a reference image and a text prompt using " "[AAD-1](https://huggingface.co/Watay/AAD-1), a one-step causal video world model " "distilled from Wan2.1-T2V-14B. \n" "[:paper: Paper](https://arxiv.org/abs/2606.03972) · " "[:computer: Code](https://github.com/AutoLab-SAI-SJTU/AAD-1) · " "[:model: Model](https://huggingface.co/Watay/AAD-1)" ) with gr.Row(): with gr.Column(scale=1): input_image = gr.Image(label="Reference Image", type="pil") prompt = gr.Textbox( label="Prompt", placeholder="Describe the motion and scene you want to generate…", lines=3, ) run_btn = gr.Button("Generate Video", variant="primary") with gr.Column(scale=1): output_video = gr.Video(label="Generated Video") with gr.Accordion("Advanced settings", open=False): num_frames = gr.Slider( label="Number of frames", minimum=17, maximum=161, step=16, value=DEFAULT_NUM_FRAMES, info="Must satisfy 16n+1. Higher = longer video but more VRAM.", ) seed = gr.Number(label="Seed", value=DEFAULT_SEED, precision=0) fps = gr.Slider( label="Output FPS", minimum=4, maximum=24, step=1, value=DEFAULT_FRAME_RATE, ) gr.Examples( examples=EXAMPLES, inputs=[input_image, prompt, num_frames, seed, fps], outputs=output_video, fn=generate_video, cache_examples=True, cache_mode="lazy", ) run_btn.click( fn=generate_video, inputs=[input_image, prompt, num_frames, seed, fps], outputs=output_video, api_name="generate", ) if __name__ == "__main__": demo.launch(mcp_server=True)