Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
22 commits
Select commit Hold shift + click to select a range
6365601
support wan-animate-2
kelseyee Jul 30, 2026
14f482c
Merge branch 'huggingface:main' into main
kelseyee Aug 7, 2026
a54425a
Refactor Wan-Animate-2 to diffusers conventions
Aug 7, 2026
2b722ed
Free per-segment KV cache and latents at the segment boundary
yiyixuxu Aug 8, 2026
270a84e
Let `load_video` report the video's frame rate
yiyixuxu Aug 8, 2026
fc0c23b
Preprocess Wan-Animate-2 with the image/video processors
yiyixuxu Aug 8, 2026
54dabb7
Match Wan-Animate-2 preprocessing to the reference implementation
yiyixuxu Aug 10, 2026
918c239
Add Wan-Animate-2 modular pipeline
yiyixuxu Aug 11, 2026
cc5c239
Repack Wan-Animate-2 modular blocks into canonical standalone steps
yiyixuxu Aug 11, 2026
e0291f8
Refine Wan-Animate-2 modular blocks: in-loop VAE encode, v1-style ten…
yiyixuxu Aug 11, 2026
1b81e87
Add Wan-Animate-2 modular tests, fill docstrings, self-assemble block…
yiyixuxu Aug 12, 2026
b9a1586
Add Wan-Animate-2 docs (modular pipeline + transformer)
yiyixuxu Aug 12, 2026
1aa8d2f
Merge branch 'main' into animate2-refactor
yiyixuxu Aug 12, 2026
1be4ed4
Run make style / make quality
yiyixuxu Aug 12, 2026
c307609
Point Wan-Animate-2 tests at the hf-internal-testing tiny repos
yiyixuxu Aug 12, 2026
d8965f7
Merge branch 'animate2-refactor' of github.com:yiyixuxu/diffusers int…
yiyixuxu Aug 12, 2026
6dc53de
Import the video processor lazily to break a circular import
yiyixuxu Aug 12, 2026
80938b2
Default distilled sampling to 10 steps; fix docstring/signature misma…
yiyixuxu Aug 12, 2026
3d1046a
Pin Wan-Animate-2 block defaults; extend the defaults test for it
yiyixuxu Aug 12, 2026
0eca946
Document per-preset sampling defaults; testing guide for the defaults…
yiyixuxu Aug 12, 2026
612648c
Remove the standard WanAnimate2Pipeline
yiyixuxu Aug 12, 2026
e38e13d
Give the distilled Wan-Animate-2 preset its own model name
yiyixuxu Aug 12, 2026
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion .ai/testing.md
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,7 @@ Follow the style introduced in [#14113](https://github.com/huggingface/diffusers

- Location: `tests/modular_pipelines/<model>/test_modular_pipeline_<model>.py` (one test class per blockset / pipeline variant).
- Subclass `ModularPipelineTesterMixin` (from `..test_modular_pipelines_common`) — it runs the pipeline end-to-end (call signature, batch consistency, float16, device placement) against a tiny checkpoint.
- Set `pipeline_class`, `pipeline_blocks_class`, `pretrained_model_name_or_path`, `params` / `batch_params`, and implement `get_dummy_inputs(seed=0)`. Set `expected_workflow_blocks` to pin the block name → class ordering per workflow.
- Set `pipeline_class`, `pipeline_blocks_class`, `pretrained_model_name_or_path`, `params` / `batch_params`, and implement `get_dummy_inputs(seed=0)`. Set `expected_workflow_blocks` to pin the block name → class ordering per workflow (only for blocksets with a `_workflow_map` — with a single workflow the list would just restate the class definition), and `expected_workflow_defaults` to pin each workflow's components, pipeline configs, and inputs — required ones by name, optional ones with their defaults. A pipeline without workflows pins its full blockset under the `None` key. An optional `component_configs` entry pins config values of `from_config` components against their creating spec (e.g. the guider scale that tells a base and a distilled preset apart); pretrained components take their config from the repo, so there is nothing block-level to pin.
- `pretrained_model_name_or_path` is a tiny repo with real components (tiny transformer, real scheduler / VAE / tokenizer configs). Develop against a personal repo; tiny repos ultimately live under `hf-internal-testing/` — not merge-blocking, a maintainer moves it before or after merge.
- **The tiny repo must mirror the real checkpoint's shape** — same index file type, same pipeline-level config keys, a scheduler configured like the real one. A fixture that doesn't look like the published repos tests a loading/config path no user will ever hit, while the path users *do* hit stays uncovered. If the model ships variants with different configs (base/distilled, different schedules), make one tiny repo and test class per variant — see the flux2 klein base/distilled split.
- **Bespoke tests go on the tester class as methods**, not as module-level functions — the mixin is pytest-style, so fixtures (`tmp_path`, `pytest.raises`, parametrize) all work in methods.
Expand Down
4 changes: 4 additions & 0 deletions docs/source/en/_toctree.yml
Original file line number Diff line number Diff line change
Expand Up @@ -405,6 +405,8 @@
title: Transformer2DModel
- local: api/models/transformer_temporal
title: TransformerTemporalModel
- local: api/models/wan_animate_2_transformer_3d
title: WanAnimate2Transformer3DModel
- local: api/models/wan_animate_transformer_3d
title: WanAnimateTransformer3DModel
- local: api/models/wan_transformer_3d
Expand Down Expand Up @@ -711,6 +713,8 @@
title: Stable Video Diffusion
- local: api/pipelines/wan
title: Wan
- local: api/pipelines/wan_animate_2
title: Wan-Animate-2
title: Video
title: Pipelines
- sections:
Expand Down
30 changes: 30 additions & 0 deletions docs/source/en/api/models/wan_animate_2_transformer_3d.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,30 @@
<!-- Copyright 2026 The HuggingFace Team. All rights reserved.

Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
the License. You may obtain a copy of the License at

http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
specific language governing permissions and limitations under the License. -->

# WanAnimate2Transformer3DModel

A Diffusion Transformer model for 3D video-like data used in [Wan-Animate-2](https://github.com/Wan-Video/Wan2.2) by the Alibaba Wan Team. It animates a character image with the motion of a driving video through an in-context reference mechanism: each segment first runs a reference pass (`kv_cache_mode="extract"`) that caches every layer's reference K/V, then the denoising passes (`kv_cache_mode="cached"`) attend jointly over the generation tokens and the cached reference tokens through a flex `BlockMask`.

The model can be loaded with the following code snippet.

```python
from diffusers import WanAnimate2Transformer3DModel

transformer = WanAnimate2Transformer3DModel.from_pretrained("Wan-AI/Wan2.2-Animate-2-14B-Diffusers", subfolder="transformer", dtype=torch.bfloat16)
```

## WanAnimate2Transformer3DModel

[[autodoc]] WanAnimate2Transformer3DModel

## Transformer2DModelOutput

[[autodoc]] models.modeling_outputs.Transformer2DModelOutput
76 changes: 76 additions & 0 deletions docs/source/en/api/pipelines/wan_animate_2.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,76 @@
<!-- Copyright 2026 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License. -->

# Wan-Animate-2

[Wan-Animate-2](https://github.com/Wan-Video/Wan2.2) by the Alibaba Wan Team animates a reference character image with the motion of a driving video. The driving video is processed in fixed-length segments: each segment runs a reference-extraction pass that caches the driving segment's K/V in every transformer layer, denoises against that cache, and is decoded inside the loop because the next segment conditions on the previous segment's decoded tail frames.

Two presets are available: the base checkpoint samples with classifier-free guidance, and the distilled checkpoint samples in few steps without it (its guider is pinned to `guidance_scale=1.0`).

```python
import torch
from diffusers import ModularPipeline
from diffusers.utils import export_to_video, load_image, load_video

pipe = ModularPipeline.from_pretrained("Wan-AI/Wan2.2-Animate-2-14B-Diffusers")
pipe.load_components(dtype=torch.bfloat16)

# The transformer weights and the per-segment reference KV cache do not co-reside on one 80 GB
# card at the default resolution, so stream the transformer's blocks. The in-context attention
# runs on the flex backend; compiling fuses it.
from diffusers.hooks import apply_group_offloading

Comment on lines +32 to +33

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
from diffusers.hooks import apply_group_offloading

apply_group_offloading(
pipe.transformer,
onload_device=torch.device("cuda"),
offload_device=torch.device("cpu"),
offload_type="block_level",
use_stream=True,
)
Comment on lines +34 to +40

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
apply_group_offloading(
pipe.transformer,
onload_device=torch.device("cuda"),
offload_device=torch.device("cpu"),
offload_type="block_level",
use_stream=True,
)
pipe.transformer.enable_group_offloading(
pipe.transformer,
onload_device=torch.device("cuda"),
offload_device=torch.device("cpu"),
offload_type="block_level",
use_stream=True,
)

pipe.text_encoder.to("cuda")
pipe.image_encoder.to("cuda")
pipe.vae.to("cuda")
pipe.transformer.compile_repeated_blocks(fullgraph=False)

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Curious: what causes the graph break?


driving_video, driving_video_fps = load_video("driving.mp4", return_fps=True)

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Let's use something that is directly loadable. We don't know where "driving.mp4" is coming from.


videos = pipe(
image=load_image("reference.png"),

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Same as above.

driving_video=driving_video,
driving_video_fps=driving_video_fps,
prompt="A cat in a blue uniform, white background",
output="videos",
)
export_to_video(videos[0], "output.mp4", fps=24)
```

For the distilled checkpoint, load `Wan-AI/Wan2.2-Animate-2-14B-Distilled-Diffusers` the same way — nothing else changes. Each preset carries its own sampling defaults (40 steps for the base checkpoint, 10 for the distilled one), and no `guidance_scale` argument exists anywhere: guidance is owned by the pipeline's guider component (classifier-free guidance at 3.0 for the base preset, disabled for the distilled one).

`height` and `width` (defaults 800 and 640) set the target *area* of the generated video; the actual frame size keeps the reference image's aspect ratio, and the driving frames are letterboxed to it. Inputs that already sit at the target letterbox size pass through the preprocessing untouched, so preprocessing can also be done entirely outside the pipeline.

## WanAnimate2ModularPipeline

[[autodoc]] WanAnimate2ModularPipeline

## WanAnimate2DistilledModularPipeline

[[autodoc]] WanAnimate2DistilledModularPipeline

## WanAnimate2Blocks

[[autodoc]] WanAnimate2Blocks

## WanAnimate2DistilledBlocks

[[autodoc]] WanAnimate2DistilledBlocks
10 changes: 10 additions & 0 deletions src/diffusers/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -352,6 +352,7 @@
"UNetSpatioTemporalConditionModel",
"UVit2DModel",
"VQModel",
"WanAnimate2Transformer3DModel",
"WanAnimateTransformer3DModel",
"WanTransformer3DModel",
"WanVACETransformer3DModel",
Expand Down Expand Up @@ -558,6 +559,10 @@
"Wan22Image2VideoBlocks",
"Wan22Image2VideoModularPipeline",
"Wan22ModularPipeline",
"WanAnimate2Blocks",
"WanAnimate2DistilledBlocks",
"WanAnimate2DistilledModularPipeline",
"WanAnimate2ModularPipeline",
"WanBlocks",
"WanImage2VideoAutoBlocks",
"WanImage2VideoModularPipeline",
Expand Down Expand Up @@ -1204,6 +1209,7 @@
UNetSpatioTemporalConditionModel,
UVit2DModel,
VQModel,
WanAnimate2Transformer3DModel,
WanAnimateTransformer3DModel,
WanTransformer3DModel,
WanVACETransformer3DModel,
Expand Down Expand Up @@ -1390,6 +1396,10 @@
Wan22Image2VideoBlocks,
Wan22Image2VideoModularPipeline,
Wan22ModularPipeline,
WanAnimate2Blocks,
WanAnimate2DistilledBlocks,
WanAnimate2DistilledModularPipeline,
WanAnimate2ModularPipeline,
WanBlocks,
WanImage2VideoAutoBlocks,
WanImage2VideoModularPipeline,
Expand Down
5 changes: 5 additions & 0 deletions src/diffusers/loaders/single_file_model.py
Original file line number Diff line number Diff line change
Expand Up @@ -54,6 +54,7 @@
convert_sana_transformer_to_diffusers,
convert_sd3_transformer_checkpoint_to_diffusers,
convert_stable_cascade_unet_single_file_to_diffusers,
convert_wan_animate_2_transformer_to_diffusers,
convert_wan_transformer_to_diffusers,
convert_wan_vae_to_diffusers,
convert_z_image_controlnet_checkpoint_to_diffusers,
Expand Down Expand Up @@ -172,6 +173,10 @@
"checkpoint_mapping_fn": convert_wan_transformer_to_diffusers,
"default_subfolder": "transformer",
},
"WanAnimate2Transformer3DModel": {
"checkpoint_mapping_fn": convert_wan_animate_2_transformer_to_diffusers,
"default_subfolder": "transformer",
},
"AutoencoderKLWan": {
"checkpoint_mapping_fn": convert_wan_vae_to_diffusers,
"default_subfolder": "vae",
Expand Down
32 changes: 32 additions & 0 deletions src/diffusers/loaders/single_file_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -3289,6 +3289,38 @@ def reshape_bias_handler(key, state_dict):
return converted_state_dict


def convert_wan_animate_2_transformer_to_diffusers(checkpoint, **kwargs):
r"""
Converts the state dict of the Wan-Animate-2 transformer from the official checkpoint format to the diffusers
format.
"""
attention_renames = {
".q.": ".to_q.",
".k.": ".to_k.",
".v.": ".to_v.",
".o.": ".to_out.0.",
".k_img.": ".add_k_proj.",
".v_img.": ".add_v_proj.",
".norm_k_img.": ".norm_added_k.",
}

converted_state_dict = {}
for key in list(checkpoint.keys()):
new_key = key.replace("model.diffusion_model.", "")
# The official checkpoint wraps every transformer block in an in-context module the
# diffusers layout does not have: `blocks.N.block.X` -> `blocks.N.X`.
if new_key.startswith("blocks."):
new_key = new_key.replace(".block.", ".", 1)
if ".self_attn." in new_key or ".cross_attn." in new_key:
for old, new in attention_renames.items():
if old in new_key:
new_key = new_key.replace(old, new)
break
converted_state_dict[new_key] = checkpoint.pop(key)

return converted_state_dict


def convert_wan_vae_to_diffusers(checkpoint, **kwargs):
converted_state_dict = {}

Expand Down
2 changes: 2 additions & 0 deletions src/diffusers/models/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -145,6 +145,7 @@
_import_structure["transformers.transformer_temporal"] = ["TransformerTemporalModel"]
_import_structure["transformers.transformer_wan"] = ["WanTransformer3DModel"]
_import_structure["transformers.transformer_wan_animate"] = ["WanAnimateTransformer3DModel"]
_import_structure["transformers.transformer_wan_animate_2"] = ["WanAnimate2Transformer3DModel"]
_import_structure["transformers.transformer_wan_vace"] = ["WanVACETransformer3DModel"]
_import_structure["transformers.transformer_z_image"] = ["ZImageTransformer2DModel"]
_import_structure["unets.unet_1d"] = ["UNet1DModel"]
Expand Down Expand Up @@ -285,6 +286,7 @@
T5FilmDecoder,
Transformer2DModel,
TransformerTemporalModel,
WanAnimate2Transformer3DModel,
WanAnimateTransformer3DModel,
WanTransformer3DModel,
WanVACETransformer3DModel,
Expand Down
1 change: 1 addition & 0 deletions src/diffusers/models/transformers/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -64,5 +64,6 @@
from .transformer_temporal import TransformerTemporalModel
from .transformer_wan import WanTransformer3DModel
from .transformer_wan_animate import WanAnimateTransformer3DModel
from .transformer_wan_animate_2 import WanAnimate2Transformer3DModel
from .transformer_wan_vace import WanVACETransformer3DModel
from .transformer_z_image import ZImageTransformer2DModel
Loading
Loading