support-wan-animate-2

[kelseyee]

Add Wan-Animate-2 (In-Context Attention) Pipeline

Summary

This PR adds support for Wan-Animate-2, a character animation model that uses an in-context attention mechanism with KV cache and block mask (flex_attention), to the diffusers library.

Unlike the existing Wan-Animate v1 (which uses a motion encoder + face encoder), Wan-Animate-2 directly consumes driving video latents via a two-phase forward (reference encoding → generation with cached KV), eliminating intermediate motion extractors.

Model Architecture

Component	Wan-Animate v1 (existing)	Wan-Animate-2 (this PR)
Conditioning	Motion encoder + face encoder + face adapter	In-context attention with driving video latents
Attention	Standard WanTransformerBlock	IncontextAttentionBlock (forward_ref + forward_gen)
KV cache	No	Yes (reference K/V cached, reused during generation)
Block mask	No	Yes (flex_attention + create_block_mask)
Score mod	No	Yes (log_scale weighting for reference attention)
in_channels	36	36 (same)

Two-phase forward

1. forward_ref: Encodes the driving video, caches K/V per layer (40 layers × K+V)
2. forward_gen: Generates video using cached K/V + block mask for frame-level sparse in-context attention, with score_mod (log_scale) for reference attention weighting

Key components (copied verbatim from original)

• create_mask: Block mask creation logic for frame-level sparse attention
• rope_params / rope_apply: RoPE computation (preserves float64 on CUDA)
• flash_attention / flex_attention: Attention functions (with lazy torch.compile)
• _score_mod_impl: log_scale score modification

New files

File	Description
src/diffusers/models/transformers/transformer_wan_animate_2.py	WanAnimate2Transformer3DModel — transformer model (~95% verbatim copy from original, 5% diffusers shell adaptation)
src/diffusers/pipelines/wan/pipeline_wan_animate_2.py	WanAnimate2Pipeline — standard pipeline with segment-based generation, FPS resampling, letterbox resize
src/diffusers/modular_pipelines/wan/modular_blocks_wan_animate_2.py	WanAnimate2Blocks / WanAnimate2ModularPipeline — modular pipeline blocks
tests/models/transformers/test_models_transformer_wan_animate_2.py	Model-level tests (two-phase forward + weight mapping)
tests/pipelines/wan/test_wan_animate_2.py	Pipeline-level tests

Modified files

File	Change
src/diffusers/__init__.py	Register WanAnimate2Transformer3DModel, WanAnimate2Pipeline
src/diffusers/models/__init__.py	Register transformer model
src/diffusers/models/transformers/__init__.py	Lazy import for transformer
src/diffusers/pipelines/__init__.py	Register pipeline
src/diffusers/pipelines/wan/__init__.py	Register pipeline in wan module
src/diffusers/modular_pipelines/wan/__init__.py	Register modular blocks + pipeline
src/diffusers/modular_pipelines/wan/modular_pipeline.py	Add WanAnimate2ModularPipeline class
src/diffusers/loaders/single_file_model.py	Register single-file loading
src/diffusers/loaders/single_file_utils.py	Add convert_wan_animate_2_transformer_to_diffusers
src/diffusers/pipelines/pipeline_utils.py	Add try/except for huggingface_hub compatibility (optional APIs)
src/diffusers/utils/dummy_pt_objects.py	Auto-generated
src/diffusers/utils/dummy_torch_and_transformers_objects.py	Auto-generated

Usage

Basic inference

import torch
from diffusers import WanAnimate2Pipeline
from diffusers.utils import export_to_video, load_image

pipe = WanAnimate2Pipeline.from_pretrained(
    "Wan2.2-Animate-2-14B-Diffusers", torch_dtype=torch.bfloat16
).to("cuda")

output = pipe(
    image=load_image("reference.png"),
    driving_video="driving.mp4",
    prompt="A person in a red shirt",
    height=800,
    width=640,
    num_inference_steps=40,
)

export_to_video(output.frames[0], "output.mp4", fps=24)

Distilled model (10 steps, no CFG)

pipe = WanAnimate2Pipeline.from_pretrained(
    "Wan2.2-Animate-2-14B-Distilled-Diffusers", torch_dtype=torch.bfloat16
).to("cuda")

output = pipe(
    image=load_image("reference.png"),
    driving_video="driving.mp4",
    prompt="...",
    num_inference_steps=10,
    guidance_scale=1.0,        # no classifier-free guidance
    flow_solver="euler",       # Euler scheduler for distilled model
)

export_to_video(output.frames[0], "output.mp4", fps=24)

Key parameters

Parameter	Default	Description
image	required	Reference character image
driving_video	required	Driving video path (str) or list of PIL images
prompt	required	Text description of character appearance
height / width	800 / 640	Output resolution (letterbox resized to match aspect ratio)
clip_len	81	Frames per segment
num_inference_steps	40	Denoising steps (10 for distilled)
guidance_scale	3.0	CFG scale (1.0 = no CFG, for distilled)
flow_solver	"dpm"	Scheduler type ("dpm" or "euler")
fps	24	Output FPS (driving video resampled to this FPS)
sample_shift	5.0	Sigma shift for flow matching

Design decisions

1. Maximum code copying: ~95% of transformer computation logic is verbatim copied from the original wan_animate_2_model.py. Only the outer shell (class inheritance, @register_to_config, class attributes, gradient checkpointing) is adapted.

2. Preserves original behavior:

   • float64 for RoPE (matches transformer_wan.py conditional dtype pattern)
   • flex_attention with block_mask + score_mod called directly (not through dispatch_attention_fn)
   • Lazy torch.compile for flex_attention (compiled on first call, not at import)

3. Pipeline matches original:

   • FPS resampling via decord + get_frame_indices
   • Letterbox resize via resize_by_area (keep aspect ratio + black padding)
   • CLIP encoding via direct bicubic resize to 224×224 (not CLIPImageProcessor)
   • Post-generation padding removal

4. Shared components: VAE (AutoencoderKLWan), text encoder (UMT5EncoderModel), CLIP (CLIPVisionModel), and scheduler (DPMSolverMultistepScheduler with flow_prediction) are the same as Wan2.1.

Memory requirements

Resolution	Transformer	KV cache	Total (with CPU offload)
480×320	28GB	21GB	~55GB
640×480	28GB	26GB	~63GB
800×640	28GB	35GB	~72GB

Use pipe.enable_model_cpu_offload() to offload T5/CLIP/VAE to CPU when not in use.

Checklist

• Transformer model (WanAnimate2Transformer3DModel)
• Standard pipeline (WanAnimate2Pipeline)
• Modular pipeline blocks (WanAnimate2Blocks, WanAnimate2ModularPipeline)
• Single-file weight loading (from_single_file)
• All __init__.py registrations
• Model-level and pipeline-level tests
• make style passed (ruff check + format)
• make fix-copies passed
• End-to-end inference verified (480×320, 640×480)
• make style re-run after latest changes
• Full test suite passes on GPU

Known limitations

1. flex_attention requires torch.compile: Without compilation, falls back to dense (math) implementation which OOMs on large resolutions. Lazy compile is used to avoid import-time compilation.

2. flash_attn is a hard dependency: The flash_attention function requires flash_attn package. Not guarded with is_available because it's essential for the model's attention computation.

3. enable_model_cpu_offload() may break KV cache: The accelerate hooks may interfere with the KV cache dict. Use .to("cuda") for reliable operation, or ensure sufficient GPU memory.

4. Context parallel not included: The original model supports FSDP/tensor parallel via wanxiang.ops. This is not included in the initial PR. Multi-GPU inference requires manual FSDP setup.

@yiyixuxu

[github-actions]

Hi @kelseyee, thanks for the PR! It does not appear to link an issue it fixes. If this PR addresses an existing issue, please add a closing keyword (e.g. Fixes #1234) to the PR description so the issue is linked. See the contribution guide for more details. If this PR intentionally does not fix a tracked issue, a maintainer can add the no-issue-needed label to silence this reminder.