app.py

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