diff --git a/.ai/testing.md b/.ai/testing.md index 24d39da3bd68..03df2520a42c 100644 --- a/.ai/testing.md +++ b/.ai/testing.md @@ -36,7 +36,7 @@ Follow the style introduced in [#14113](https://github.com/huggingface/diffusers - Location: `tests/modular_pipelines//test_modular_pipeline_.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. diff --git a/docs/source/en/_toctree.yml b/docs/source/en/_toctree.yml index c90b68b2b54c..9b65cd97bd3e 100644 --- a/docs/source/en/_toctree.yml +++ b/docs/source/en/_toctree.yml @@ -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 @@ -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: diff --git a/docs/source/en/api/models/wan_animate_2_transformer_3d.md b/docs/source/en/api/models/wan_animate_2_transformer_3d.md new file mode 100644 index 000000000000..d21e2b70bd62 --- /dev/null +++ b/docs/source/en/api/models/wan_animate_2_transformer_3d.md @@ -0,0 +1,30 @@ + + +# 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 diff --git a/docs/source/en/api/pipelines/wan_animate_2.md b/docs/source/en/api/pipelines/wan_animate_2.md new file mode 100644 index 000000000000..7e23c8179ec3 --- /dev/null +++ b/docs/source/en/api/pipelines/wan_animate_2.md @@ -0,0 +1,76 @@ + + +# 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 + +apply_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) + +driving_video, driving_video_fps = load_video("driving.mp4", return_fps=True) + +videos = pipe( + image=load_image("reference.png"), + 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 diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py index 03ed1fcb34ae..66cfa442908b 100644 --- a/src/diffusers/__init__.py +++ b/src/diffusers/__init__.py @@ -352,6 +352,7 @@ "UNetSpatioTemporalConditionModel", "UVit2DModel", "VQModel", + "WanAnimate2Transformer3DModel", "WanAnimateTransformer3DModel", "WanTransformer3DModel", "WanVACETransformer3DModel", @@ -558,6 +559,10 @@ "Wan22Image2VideoBlocks", "Wan22Image2VideoModularPipeline", "Wan22ModularPipeline", + "WanAnimate2Blocks", + "WanAnimate2DistilledBlocks", + "WanAnimate2DistilledModularPipeline", + "WanAnimate2ModularPipeline", "WanBlocks", "WanImage2VideoAutoBlocks", "WanImage2VideoModularPipeline", @@ -1204,6 +1209,7 @@ UNetSpatioTemporalConditionModel, UVit2DModel, VQModel, + WanAnimate2Transformer3DModel, WanAnimateTransformer3DModel, WanTransformer3DModel, WanVACETransformer3DModel, @@ -1390,6 +1396,10 @@ Wan22Image2VideoBlocks, Wan22Image2VideoModularPipeline, Wan22ModularPipeline, + WanAnimate2Blocks, + WanAnimate2DistilledBlocks, + WanAnimate2DistilledModularPipeline, + WanAnimate2ModularPipeline, WanBlocks, WanImage2VideoAutoBlocks, WanImage2VideoModularPipeline, diff --git a/src/diffusers/loaders/single_file_model.py b/src/diffusers/loaders/single_file_model.py index 56770fd9b6c3..a07657159d36 100644 --- a/src/diffusers/loaders/single_file_model.py +++ b/src/diffusers/loaders/single_file_model.py @@ -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, @@ -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", diff --git a/src/diffusers/loaders/single_file_utils.py b/src/diffusers/loaders/single_file_utils.py index 296f32f891f0..b5c6846ba17c 100644 --- a/src/diffusers/loaders/single_file_utils.py +++ b/src/diffusers/loaders/single_file_utils.py @@ -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 = {} diff --git a/src/diffusers/models/__init__.py b/src/diffusers/models/__init__.py index 6d31515eed8f..f9f40899f4a9 100755 --- a/src/diffusers/models/__init__.py +++ b/src/diffusers/models/__init__.py @@ -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"] @@ -285,6 +286,7 @@ T5FilmDecoder, Transformer2DModel, TransformerTemporalModel, + WanAnimate2Transformer3DModel, WanAnimateTransformer3DModel, WanTransformer3DModel, WanVACETransformer3DModel, diff --git a/src/diffusers/models/transformers/__init__.py b/src/diffusers/models/transformers/__init__.py index 7a1213639e3d..2333acc06762 100755 --- a/src/diffusers/models/transformers/__init__.py +++ b/src/diffusers/models/transformers/__init__.py @@ -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 diff --git a/src/diffusers/models/transformers/transformer_wan_animate_2.py b/src/diffusers/models/transformers/transformer_wan_animate_2.py new file mode 100644 index 000000000000..9259eae38342 --- /dev/null +++ b/src/diffusers/models/transformers/transformer_wan_animate_2.py @@ -0,0 +1,949 @@ +# 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. + +import math + +import torch +import torch.nn as nn +from torch.nn.attention.flex_attention import BlockMask, create_block_mask + +from ...configuration_utils import ConfigMixin, register_to_config +from ...loaders import FromOriginalModelMixin, PeftAdapterMixin +from ..attention import AttentionMixin, AttentionModuleMixin +from ..attention_dispatch import AttentionBackendName, dispatch_attention_fn +from ..embeddings import Timesteps +from ..modeling_outputs import Transformer2DModelOutput +from ..modeling_utils import ModelMixin +from ..normalization import FP32LayerNorm + + +def rope_params(max_seq_len, dim, theta=10000, offset=0): + assert dim % 2 == 0 + freqs = torch.outer( + torch.arange(max_seq_len) + offset, + 1.0 / torch.pow(theta, torch.arange(0, dim, 2).to(torch.float64).div(dim)), + ) + freqs = torch.polar(torch.ones_like(freqs), freqs) + return freqs + + +def rope_apply(x, grid_sizes, freqs, time_stride=1): + n, c = x.size(2), x.size(3) // 2 + + # split freqs + freqs = freqs.split([c - 2 * (c // 3), c // 3, c // 3], dim=1) + + # loop over samples + output = [] + for i, (f, h, w) in enumerate(grid_sizes.tolist()): + seq_len = f * h * w + + # precompute multipliers + x_i = torch.view_as_complex(x[i, :seq_len].to(torch.float64).reshape(seq_len, n, -1, 2)) + freqs_i = torch.cat( + [ + freqs[0][: f * time_stride : time_stride].view(f, 1, 1, -1).expand(f, h, w, -1), + freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1), + freqs[2][:w].view(1, 1, w, -1).expand(f, h, w, -1), + ], + dim=-1, + ).reshape(seq_len, 1, -1) + + # apply rotary embedding + x_i = torch.view_as_real(x_i * freqs_i).flatten(2) + x_i = torch.cat([x_i, x[i, seq_len:]]) + + # append to collection + output.append(x_i) + return torch.stack(output).float() + + +def pad_freqs(original_tensor, target_len): + seq_len, s1, s2 = original_tensor.shape + pad_size = target_len - seq_len + padding_tensor = torch.ones( + pad_size, + s1, + s2, + dtype=original_tensor.dtype, + device=original_tensor.device, + ) + padded_tensor = torch.cat([original_tensor, padding_tensor], dim=0) + return padded_tensor + + +def _get_qkv_projections(attn, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor | None): + # encoder_hidden_states is only passed for cross-attention + if encoder_hidden_states is None: + encoder_hidden_states = hidden_states + + if attn.fused_projections: + if not attn.is_cross_attention: + # In self-attention layers, we can fuse the entire QKV projection into a single linear + query, key, value = attn.to_qkv(hidden_states).chunk(3, dim=-1) + else: + # In cross-attention layers, we can only fuse the KV projections into a single linear + query = attn.to_q(hidden_states) + key, value = attn.to_kv(encoder_hidden_states).chunk(2, dim=-1) + else: + query = attn.to_q(hidden_states) + key = attn.to_k(encoder_hidden_states) + value = attn.to_v(encoder_hidden_states) + return query, key, value + + +def _get_added_kv_projections(attn, encoder_hidden_states_img: torch.Tensor): + if attn.fused_projections: + key_img, value_img = attn.to_added_kv(encoder_hidden_states_img).chunk(2, dim=-1) + else: + key_img = attn.add_k_proj(encoder_hidden_states_img) + value_img = attn.add_v_proj(encoder_hidden_states_img) + return key_img, value_img + + +class WanAnimate2KVLayerCache: + """Per-layer K/V cache for the reference tokens. + + Holds the *pre-RoPE* projections: the generation pass re-applies rotary embeddings to the reference keys using the + reference grid and the `refer_offset_*` offsets. Tensor format: `(batch_size, seq_len, num_heads, head_dim)`. + """ + + def __init__(self): + self.key: torch.Tensor | None = None + self.value: torch.Tensor | None = None + + def store(self, key: torch.Tensor, value: torch.Tensor): + self.key = key + self.value = value + + def get(self) -> tuple[torch.Tensor, torch.Tensor]: + if self.key is None: + raise RuntimeError("The KV cache is empty. Run the reference pass before the generation pass.") + return self.key, self.value + + def clear(self): + self.key = None + self.value = None + + +class WanAnimate2KVCache: + """Container holding one [`WanAnimate2KVLayerCache`] per transformer layer.""" + + def __init__(self, num_layers: int): + self.layer_caches = [WanAnimate2KVLayerCache() for _ in range(num_layers)] + + def get(self, layer_idx: int) -> WanAnimate2KVLayerCache: + return self.layer_caches[layer_idx] + + def clear(self): + for cache in self.layer_caches: + cache.clear() + + +class WanAnimate2AttnProcessor: + r""" + Self-attention for the Wan-Animate-2 in-context reference mechanism. + + With `kv_cache_mode="extract"` (the reference pass) this is dense self-attention over the reference tokens; the + projected K/V are written to `kv_cache` before rotary embeddings are applied. + + With `kv_cache_mode="cached"` (the generation pass) the generation tokens and the cached reference tokens are + packed into a 128-aligned `[generation | reference]` buffer and attended through a flex `BlockMask`, so each + generation frame attends to every generation token plus the reference tokens at the same frame index. Because the + pattern is expressed as a `BlockMask`, this path runs on the `flex` backend only. + """ + + _attention_backend = None + _parallel_config = None + + def __call__( + self, + attn: "WanAnimate2Attention", + hidden_states: torch.Tensor, + rotary_emb: torch.Tensor, + grid_sizes: torch.Tensor, + kv_cache: WanAnimate2KVLayerCache, + kv_cache_mode: str, + rope_stride: int = 1, + reference_rotary_emb: torch.Tensor | None = None, + reference_grid_sizes: torch.Tensor | None = None, + reference_rope_stride: int = 1, + attention_mask: BlockMask | None = None, + origin_latent_frames: int | None = None, + origin_latent_hw: int | None = None, + ) -> torch.Tensor: + query, key, value = _get_qkv_projections(attn, hidden_states, None) + + query = attn.norm_q(query) + key = attn.norm_k(key) + + query = query.unflatten(2, (attn.heads, -1)) + key = key.unflatten(2, (attn.heads, -1)) + value = value.unflatten(2, (attn.heads, -1)) + + if kv_cache_mode == "extract": + kv_cache.store(key, value) + + # `rope_apply` computes in float64 and returns float32; attention runs in the model dtype. + query = rope_apply(query, grid_sizes, rotary_emb, rope_stride).type_as(value) + key = rope_apply(key, grid_sizes, rotary_emb, rope_stride).type_as(value) + + hidden_states = dispatch_attention_fn( + query, + key, + value, + attn_mask=None, + backend=self._attention_backend, + parallel_config=self._parallel_config, + ) + elif kv_cache_mode == "cached": + query = rope_apply(query, grid_sizes, rotary_emb, rope_stride).type_as(value) + key = rope_apply(key, grid_sizes, rotary_emb, rope_stride).type_as(value) + + key_ref, value_ref = kv_cache.get() + key_ref = rope_apply(key_ref, reference_grid_sizes, reference_rotary_emb, reference_rope_stride).type_as( + value + ) + + frames, height, width = grid_sizes[0].tolist() + ref_frames, ref_height, ref_width = reference_grid_sizes[0].tolist() + hw, ref_hw = height * width, ref_height * ref_width + valid_length, ref_valid_length = frames * hw, ref_frames * ref_hw + + batch_size, _, heads, head_dim = query.shape + + # The block mask is built once for the full video resolution, so this segment is + # scattered into a buffer of that size. Both segments are padded to a multiple of + # 128 to match the block mask's granularity. + packed_length = math.ceil((origin_latent_frames + 1) * origin_latent_hw / 128) * 128 + packed_ref_length = math.ceil(origin_latent_frames * origin_latent_hw / 128) * 128 + + query_packed = query.new_zeros(batch_size, packed_length, heads, head_dim) + key_packed = key.new_zeros(batch_size, packed_length + packed_ref_length, heads, head_dim) + value_packed = value.new_zeros(batch_size, packed_length + packed_ref_length, heads, head_dim) + + generation = slice(0, frames * origin_latent_hw) + query_packed[:, generation].view(batch_size, frames, origin_latent_hw, heads, head_dim)[:, :, :hw] = query[ + :, :valid_length + ].view(batch_size, frames, hw, heads, head_dim) + key_packed[:, generation].view(batch_size, frames, origin_latent_hw, heads, head_dim)[:, :, :hw] = key[ + :, :valid_length + ].view(batch_size, frames, hw, heads, head_dim) + value_packed[:, generation].view(batch_size, frames, origin_latent_hw, heads, head_dim)[:, :, :hw] = value[ + :, :valid_length + ].view(batch_size, frames, hw, heads, head_dim) + + reference = slice(packed_length, packed_length + ref_frames * origin_latent_hw) + key_packed[:, reference].view(batch_size, ref_frames, origin_latent_hw, heads, head_dim)[:, :, :ref_hw] = ( + key_ref[:, :ref_valid_length].view(batch_size, ref_frames, ref_hw, heads, head_dim) + ) + value_packed[:, reference].view(batch_size, ref_frames, origin_latent_hw, heads, head_dim)[ + :, :, :ref_hw + ] = value_ref[:, :ref_valid_length].view(batch_size, ref_frames, ref_hw, heads, head_dim) + + hidden_states = dispatch_attention_fn( + query_packed, + key_packed, + value_packed, + attn_mask=attention_mask, + backend=AttentionBackendName.FLEX, + parallel_config=self._parallel_config, + ) + + hidden_states = ( + hidden_states[:, generation] + .view(batch_size, frames, origin_latent_hw, heads, head_dim)[:, :, :hw] + .reshape(batch_size, valid_length, heads, head_dim) + ) + # Padded query positions are not attended and pass through unchanged. + hidden_states = torch.cat([hidden_states, query[:, valid_length:]], dim=1) + else: + raise ValueError(f"`kv_cache_mode` must be either 'extract' or 'cached', got {kv_cache_mode}.") + + hidden_states = hidden_states.flatten(2, 3).type_as(query) + hidden_states = attn.to_out[0](hidden_states) + hidden_states = attn.to_out[1](hidden_states) + return hidden_states + + +class WanAnimate2CrossAttnProcessor: + r""" + Cross-attention to the text embeddings, plus an additive branch over the CLIP image embeddings. + + The two token streams are passed as separate arguments rather than sliced out of one concatenated tensor, so the + CLIP token count does not have to be hardcoded. + """ + + _attention_backend = None + _parallel_config = None + + def __call__( + self, + attn: "WanAnimate2Attention", + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + encoder_hidden_states_image: torch.Tensor | None = None, + ) -> torch.Tensor: + query, key, value = _get_qkv_projections(attn, hidden_states, encoder_hidden_states) + + query = attn.norm_q(query) + key = attn.norm_k(key) + + query = query.unflatten(2, (attn.heads, -1)) + key = key.unflatten(2, (attn.heads, -1)) + value = value.unflatten(2, (attn.heads, -1)) + + hidden_states = dispatch_attention_fn( + query, + key, + value, + attn_mask=None, + backend=self._attention_backend, + parallel_config=self._parallel_config, + ) + hidden_states = hidden_states.flatten(2, 3).type_as(query) + + if encoder_hidden_states_image is not None: + key_image, value_image = _get_added_kv_projections(attn, encoder_hidden_states_image) + key_image = attn.norm_added_k(key_image) + + key_image = key_image.unflatten(2, (attn.heads, -1)) + value_image = value_image.unflatten(2, (attn.heads, -1)) + + hidden_states_image = dispatch_attention_fn( + query, + key_image, + value_image, + attn_mask=None, + backend=self._attention_backend, + parallel_config=self._parallel_config, + ) + hidden_states = hidden_states + hidden_states_image.flatten(2, 3).type_as(query) + + hidden_states = attn.to_out[0](hidden_states) + hidden_states = attn.to_out[1](hidden_states) + return hidden_states + + +class WanAnimate2Attention(torch.nn.Module, AttentionModuleMixin): + _default_processor_cls = WanAnimate2AttnProcessor + _available_processors = [WanAnimate2AttnProcessor, WanAnimate2CrossAttnProcessor] + + def __init__( + self, + dim: int, + heads: int, + eps: float = 1e-6, + dropout: float = 0.0, + added_kv_proj_dim: int | None = None, + processor=None, + is_cross_attention: bool = False, + ): + super().__init__() + + self.heads = heads + self.added_kv_proj_dim = added_kv_proj_dim + self.is_cross_attention = is_cross_attention + self.use_bias = True + + self.to_q = torch.nn.Linear(dim, dim, bias=True) + self.to_k = torch.nn.Linear(dim, dim, bias=True) + self.to_v = torch.nn.Linear(dim, dim, bias=True) + self.to_out = torch.nn.ModuleList([torch.nn.Linear(dim, dim, bias=True), torch.nn.Dropout(dropout)]) + self.norm_q = torch.nn.RMSNorm(dim, eps=eps, elementwise_affine=True) + self.norm_k = torch.nn.RMSNorm(dim, eps=eps, elementwise_affine=True) + + self.add_k_proj = self.add_v_proj = None + if added_kv_proj_dim is not None: + self.add_k_proj = torch.nn.Linear(added_kv_proj_dim, dim, bias=True) + self.add_v_proj = torch.nn.Linear(added_kv_proj_dim, dim, bias=True) + self.norm_added_k = torch.nn.RMSNorm(dim, eps=eps, elementwise_affine=True) + + self.set_processor(processor if processor is not None else self._default_processor_cls()) + + # Copied from diffusers.models.transformers.transformer_wan.WanAttention.fuse_projections + def fuse_projections(self): + if getattr(self, "fused_projections", False): + return + + if not self.is_cross_attention: + concatenated_weights = torch.cat([self.to_q.weight.data, self.to_k.weight.data, self.to_v.weight.data]) + concatenated_bias = torch.cat([self.to_q.bias.data, self.to_k.bias.data, self.to_v.bias.data]) + out_features, in_features = concatenated_weights.shape + with torch.device("meta"): + self.to_qkv = nn.Linear(in_features, out_features, bias=True) + self.to_qkv.load_state_dict( + {"weight": concatenated_weights, "bias": concatenated_bias}, strict=True, assign=True + ) + else: + concatenated_weights = torch.cat([self.to_k.weight.data, self.to_v.weight.data]) + concatenated_bias = torch.cat([self.to_k.bias.data, self.to_v.bias.data]) + out_features, in_features = concatenated_weights.shape + with torch.device("meta"): + self.to_kv = nn.Linear(in_features, out_features, bias=True) + self.to_kv.load_state_dict( + {"weight": concatenated_weights, "bias": concatenated_bias}, strict=True, assign=True + ) + + if self.added_kv_proj_dim is not None: + concatenated_weights = torch.cat([self.add_k_proj.weight.data, self.add_v_proj.weight.data]) + concatenated_bias = torch.cat([self.add_k_proj.bias.data, self.add_v_proj.bias.data]) + out_features, in_features = concatenated_weights.shape + with torch.device("meta"): + self.to_added_kv = nn.Linear(in_features, out_features, bias=True) + self.to_added_kv.load_state_dict( + {"weight": concatenated_weights, "bias": concatenated_bias}, strict=True, assign=True + ) + + self.fused_projections = True + + @torch.no_grad() + # Copied from diffusers.models.transformers.transformer_wan.WanAttention.unfuse_projections + def unfuse_projections(self): + if not getattr(self, "fused_projections", False): + return + + if hasattr(self, "to_qkv"): + delattr(self, "to_qkv") + if hasattr(self, "to_kv"): + delattr(self, "to_kv") + if hasattr(self, "to_added_kv"): + delattr(self, "to_added_kv") + + self.fused_projections = False + + def forward(self, hidden_states: torch.Tensor, **kwargs) -> torch.Tensor: + return self.processor(self, hidden_states, **kwargs) + + +class WanAnimate2TransformerBlock(nn.Module): + def __init__( + self, + dim, + ffn_dim, + num_heads, + cross_attn_norm=False, + eps=1e-6, + refer_stride=1, + use_img_emb=True, + ): + super().__init__() + self.refer_stride = refer_stride + + # 1. Self-attention + self.norm1 = FP32LayerNorm(dim, eps, elementwise_affine=False) + self.self_attn = WanAnimate2Attention( + dim=dim, + heads=num_heads, + eps=eps, + processor=WanAnimate2AttnProcessor(), + ) + + # 2. Cross-attention + self.norm3 = FP32LayerNorm(dim, eps, elementwise_affine=True) if cross_attn_norm else nn.Identity() + self.cross_attn = WanAnimate2Attention( + dim=dim, + heads=num_heads, + eps=eps, + added_kv_proj_dim=dim if use_img_emb else None, + is_cross_attention=True, + processor=WanAnimate2CrossAttnProcessor(), + ) + + # 3. Feed-forward + self.norm2 = FP32LayerNorm(dim, eps, elementwise_affine=False) + self.ffn = nn.Sequential( + nn.Linear(dim, ffn_dim), + nn.GELU(approximate="tanh"), + nn.Linear(ffn_dim, dim), + ) + + self.modulation = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5) + + def forward( + self, + hidden_states: torch.Tensor, + temb: torch.Tensor, + encoder_hidden_states: torch.Tensor, + kv_cache: WanAnimate2KVLayerCache, + kv_cache_mode: str, + rotary_emb: torch.Tensor, + grid_sizes: torch.Tensor, + encoder_hidden_states_image: torch.Tensor | None = None, + reference_rotary_emb: torch.Tensor | None = None, + reference_grid_sizes: torch.Tensor | None = None, + attention_mask: BlockMask | None = None, + origin_latent_frames: int | None = None, + origin_latent_hw: int | None = None, + ) -> torch.Tensor: + shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = (self.modulation + temb).chunk(6, dim=1) + + # 1. Self-attention. The reference tokens sit on a strided time axis, so `refer_stride` + # applies to whichever stream is the reference one in this mode. + norm_hidden_states = (self.norm1(hidden_states.float()) * (1 + scale_msa) + shift_msa).type_as(hidden_states) + attn_output = self.self_attn( + norm_hidden_states, + rotary_emb=rotary_emb, + grid_sizes=grid_sizes, + kv_cache=kv_cache, + kv_cache_mode=kv_cache_mode, + rope_stride=self.refer_stride if kv_cache_mode == "extract" else 1, + reference_rotary_emb=reference_rotary_emb, + reference_grid_sizes=reference_grid_sizes, + reference_rope_stride=self.refer_stride, + attention_mask=attention_mask, + origin_latent_frames=origin_latent_frames, + origin_latent_hw=origin_latent_hw, + ) + hidden_states = (hidden_states.float() + attn_output * gate_msa).type_as(hidden_states) + + # 2. Cross-attention + hidden_states = hidden_states + self.cross_attn( + self.norm3(hidden_states.float()).type_as(hidden_states), + encoder_hidden_states=encoder_hidden_states, + encoder_hidden_states_image=encoder_hidden_states_image, + ) + + # 3. Feed-forward + norm_hidden_states = (self.norm2(hidden_states.float()) * (1 + c_scale_msa) + c_shift_msa).type_as( + hidden_states + ) + ff_output = self.ffn(norm_hidden_states) + hidden_states = (hidden_states.float() + ff_output.float() * c_gate_msa).type_as(hidden_states) + + return hidden_states + + +class Head(nn.Module): + def __init__(self, dim, out_dim, patch_size, eps=1e-6): + super().__init__() + self.dim = dim + self.out_dim = out_dim + self.patch_size = patch_size + self.eps = eps + + # layers + out_dim = math.prod(patch_size) * out_dim + self.norm = FP32LayerNorm(dim, eps, elementwise_affine=False) + self.head = nn.Linear(dim, out_dim) + + # modulation + self.modulation = nn.Parameter(torch.randn(1, 2, dim) / dim**0.5) + + def forward(self, x, e): + shift, scale = (self.modulation + e.float().unsqueeze(1)).chunk(2, dim=1) + x = self.head((self.norm(x.float()) * (1 + scale) + shift).type_as(x)) + return x + + +class MLPProj(torch.nn.Module): + def __init__(self, in_dim, out_dim): + super().__init__() + + self.proj = torch.nn.Sequential( + torch.nn.LayerNorm(in_dim), + torch.nn.Linear(in_dim, in_dim), + torch.nn.GELU(), + torch.nn.Linear(in_dim, out_dim), + torch.nn.LayerNorm(out_dim), + ) + + def forward(self, image_embeds): + clip_extra_context_tokens = self.proj(image_embeds) + return clip_extra_context_tokens + + +class WanAnimate2Transformer3DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin, AttentionMixin): + r""" + A Transformer model for video-like data used in the Wan-Animate-2 model. + + Wan-Animate-2 uses an in-context attention mechanism with a KV cache: a reference video is first encoded + (``kv_cache_mode="extract"``) to populate a [`WanAnimate2KVCache`], then each denoising step + (``kv_cache_mode="cached"``) attends jointly over the generation tokens and the cached reference K/V through a flex + ``BlockMask``. The generation self-attention therefore runs on the ``flex`` attention backend only; every other + attention in the model works on any backend. + + Args: + patch_size (`tuple[int]`, defaults to `(1, 2, 2)`): + 3D patch dimensions for video embedding (t_patch, h_patch, w_patch). + text_len (`int`, defaults to `512`): + Fixed length for text embeddings. + in_dim (`int`, defaults to `36`): + The number of channels in the input (2 * latent_channels + 4 for mask channel). + dim (`int`, defaults to `5120`): + The number of channels in the transformer. + ffn_dim (`int`, defaults to `13824`): + Intermediate dimension in feed-forward network. + freq_dim (`int`, defaults to `256`): + Dimension for sinusoidal time embeddings. + text_dim (`int`, defaults to `4096`): + Input dimension for text embeddings. + out_dim (`int`, defaults to `16`): + The number of channels in the output. + num_heads (`int`, defaults to `40`): + The number of attention heads. + num_layers (`int`, defaults to `40`): + The number of layers of transformer blocks to use. + cross_attn_norm (`bool`, defaults to `True`): + Enable cross-attention normalization. + eps (`float`, defaults to `1e-6`): + Epsilon value for normalization layers. + use_img_emb (`bool`, defaults to `True`): + Whether to use CLIP image embedding. + refer_offset_t (`int`, defaults to `1`): + RoPE offset for the temporal dimension of the reference. + refer_offset_h (`int`, defaults to `0`): + RoPE offset for the height dimension of the reference. + refer_offset_w (`int`, defaults to `-1`): + RoPE offset for the width dimension of the reference. -1 means use the generation grid size. + refer_stride (`int`, defaults to `1`): + Stride for RoPE application on the reference. + """ + + _supports_gradient_checkpointing = True + _skip_layerwise_casting_patterns = ["patch_embedding", "img_emb", "norm"] + _no_split_modules = ["WanAnimate2TransformerBlock"] + _repeated_blocks = ["WanAnimate2TransformerBlock"] + _skip_keys = ["kv_cache"] + _keep_in_fp32_modules = [ + "time_embedding", + "time_projection", + "scale_shift_table", + "norm1", + "norm2", + "norm3", + "modulation", + ] + + @register_to_config + def __init__( + self, + patch_size: tuple = (1, 2, 2), + text_len: int = 512, + in_dim: int = 36, + dim: int = 5120, + ffn_dim: int = 13824, + freq_dim: int = 256, + text_dim: int = 4096, + out_dim: int = 16, + num_heads: int = 40, + num_layers: int = 40, + cross_attn_norm: bool = True, + eps: float = 1e-6, + use_img_emb: bool = True, + refer_offset_t: int = 1, + refer_offset_h: int = 0, + refer_offset_w: int = -1, + refer_stride: int = 1, + ): + super().__init__() + self.patch_size = patch_size + self.text_len = text_len + self.in_dim = in_dim + self.dim = dim + self.ffn_dim = ffn_dim + self.freq_dim = freq_dim + self.text_dim = text_dim + self.out_dim = out_dim + self.num_heads = num_heads + self.num_layers = num_layers + self.cross_attn_norm = cross_attn_norm + self.eps = eps + self.use_img_emb = use_img_emb + self.refer_offset_t = refer_offset_t + self.refer_offset_h = refer_offset_h + self.refer_offset_w = refer_offset_w + self.refer_stride = refer_stride + + # [Denoising Transformer] + # embeddings + self.patch_embedding = nn.Conv3d(in_dim, dim, kernel_size=patch_size, stride=patch_size) + self.text_embedding = nn.Sequential( + nn.Linear(text_dim, dim), + nn.GELU(approximate="tanh"), + nn.Linear(dim, dim), + ) + + self.timesteps_proj = Timesteps(num_channels=freq_dim, flip_sin_to_cos=True, downscale_freq_shift=0) + self.time_embedding = nn.Sequential( + nn.Linear(freq_dim, dim), + nn.SiLU(), + nn.Linear(dim, dim), + ) + self.time_projection = nn.Sequential( + nn.SiLU(), + nn.Linear(dim, dim * 6), + ) + + # blocks + self.blocks = nn.ModuleList( + [ + WanAnimate2TransformerBlock( + dim, + ffn_dim, + num_heads, + cross_attn_norm, + eps, + refer_stride, + use_img_emb=use_img_emb, + ) + for _ in range(num_layers) + ] + ) + + # head + self.head = Head(dim, out_dim, patch_size, eps) + + if use_img_emb: + self.img_emb = MLPProj(1280, dim) + + self.gradient_checkpointing = False + self.block_masks = {} + self.block_mask_grid_sizes = {} + + def create_mask(self, origin_latent_f, hw, device): + q_len = (origin_latent_f + 1) * hw + k_len = origin_latent_f * hw + + q_len_total = math.ceil(q_len / 128) * 128 + k_extra_len_total = math.ceil(k_len / 128) * 128 + k_len_total = q_len_total + k_extra_len_total + + q_limit = q_len + k_limit = k_len + q_total = q_len_total + + def attention_mask_logic(b, h, q_idx, kv_idx): + q_valid = q_idx < q_limit + is_base_attention = kv_idx < q_limit + + q_frame = q_idx // hw + is_first_part = kv_idx < q_total + + kv_frame_1 = kv_idx // hw + kv_is_valid_1 = kv_idx < q_limit + + rel_kv_idx = kv_idx - q_total + kv_frame_2 = (rel_kv_idx // hw) + 1 + kv_is_valid_2 = rel_kv_idx < k_limit + + kv_frame = torch.where(is_first_part, kv_frame_1, kv_frame_2) + kv_is_valid = torch.where(is_first_part, kv_is_valid_1, kv_is_valid_2) + + is_cond_attention = (q_frame == kv_frame) & kv_is_valid + + return q_valid & (is_base_attention | is_cond_attention) + + block_mask = create_block_mask( + attention_mask_logic, + B=None, + H=None, + Q_LEN=q_len_total, + KV_LEN=k_len_total, + device=device, + _compile=True, + ) + return block_mask + + def forward( + self, + hidden_states: list[torch.Tensor], + timestep: torch.Tensor, + encoder_hidden_states: list[torch.Tensor], + condition_latents: list[torch.Tensor], + kv_cache: WanAnimate2KVCache, + kv_cache_mode: str, + seq_len: int, + encoder_hidden_states_image: torch.Tensor | None = None, + offset_grid_sizes: torch.Tensor | None = None, + reference_grid_sizes: torch.Tensor | None = None, + origin_len: int | None = None, + origin_area: list[int] | None = None, + is_uncondtion: bool = False, + return_dict: bool = True, + ) -> Transformer2DModelOutput | tuple[list[torch.Tensor]]: + r""" + Args: + hidden_states (`list[torch.Tensor]`): + Latents for this pass — the reference latents when `kv_cache_mode="extract"`, the noisy generation + latents when `kv_cache_mode="cached"`. + timestep (`torch.Tensor`): + Denoising timestep. Ignored under `kv_cache_mode="extract"`, which uses a fixed timestep of 1. + encoder_hidden_states (`list[torch.Tensor]`): + Text embeddings for this pass. + condition_latents (`list[torch.Tensor]`): + Conditioning latents concatenated to `hidden_states` before patch embedding. + kv_cache (`WanAnimate2KVCache`): + Written under `kv_cache_mode="extract"`, read under `"cached"`. + kv_cache_mode (`str`): + `"extract"` runs the reference pass and populates `kv_cache`; `"cached"` runs a denoising step against + the cached reference tokens. + seq_len (`int`): + Token count each sample must hold after patch embedding. + encoder_hidden_states_image (`torch.Tensor`, *optional*): + CLIP image embeddings, used when the model is configured with `use_img_emb`. + offset_grid_sizes (`torch.Tensor`, *optional*): + Patch grid used to resolve any `refer_offset_*` still set to -1. Required under + `kv_cache_mode="extract"`; under `"cached"` the grid derived from `hidden_states` is used instead. Note + the two are not the same grid — the reference pass runs first and resolves the offsets from whatever it + is given here, which the pipeline sets to the *reference* grid. + reference_grid_sizes (`torch.Tensor`, *optional*): + Patch grid of the reference latents, used for the reference rotary embeddings. Required under + `kv_cache_mode="cached"`. + origin_len (`int`, *optional*): + Frame count of the full video, which the in-context block mask is built over. Required under + `kv_cache_mode="cached"`. + origin_area (`list[int]`, *optional*): + Spatial size `[height, width]` of the full video, which the in-context block mask is built over. + Required under `kv_cache_mode="cached"`. + is_uncondtion (`bool`, *optional*): + Whether this is the unconditional branch of classifier-free guidance. + return_dict (`bool`, *optional*, defaults to `True`): + Whether to return a [`~models.transformer_2d.Transformer2DModelOutput`] instead of a plain tuple. + + Returns: + [`~models.modeling_outputs.Transformer2DModelOutput`] or `tuple(list[torch.Tensor])`: + The predicted sample per input latent, unpatchified; a plain tuple if `return_dict` is `False`. + """ + if kv_cache_mode not in ("extract", "cached"): + raise ValueError(f"`kv_cache_mode` must be either 'extract' or 'cached', got {kv_cache_mode}.") + + device = self.patch_embedding.weight.device + + # 1. Patch embedding. `grid_sizes` describes whichever stream this pass is running over. + hidden_states = [torch.cat([u, v], dim=0) for u, v in zip(hidden_states, condition_latents)] + hidden_states = [self.patch_embedding(u.unsqueeze(0)) for u in hidden_states] + grid_sizes = torch.stack([torch.tensor(u.shape[2:], dtype=torch.long) for u in hidden_states]) + hidden_states = [u.flatten(2).transpose(1, 2) for u in hidden_states] + if any(u.size(1) != seq_len for u in hidden_states): + raise ValueError( + f"Each sample must hold exactly `seq_len` ({seq_len}) tokens, got " + f"{[u.size(1) for u in hidden_states]}. Self-attention here is either dense and unmasked or driven " + f"by a block mask built from the grid, so a padded sequence would not line up." + ) + hidden_states = torch.cat(hidden_states) + + assert (self.dim % self.num_heads) == 0 and (self.dim // self.num_heads) % 2 == 0 + d = self.dim // self.num_heads + + # 2. Rotary embeddings for the reference stream. The offsets place the reference tokens + # past the generation grid, so they are resolved from the generation grid either way. + offset_grid_sizes = offset_grid_sizes if kv_cache_mode == "extract" else grid_sizes + if self.refer_offset_t < 0: + self.refer_offset_t = offset_grid_sizes[0][0].item() + if self.refer_offset_h < 0: + self.refer_offset_h = offset_grid_sizes[0][1].item() + if self.refer_offset_w < 0: + self.refer_offset_w = offset_grid_sizes[0][2].item() + + self.freqs_ref = torch.cat( + [ + rope_params(512, d - 4 * (d // 6), offset=self.refer_offset_t), + rope_params(512, 2 * (d // 6), offset=self.refer_offset_h), + rope_params(512, 2 * (d // 6), offset=self.refer_offset_w), + ], + dim=1, + ) + if self.freqs_ref.device != device: + self.freqs_ref = self.freqs_ref.to(device) + + # 3. Time and context embeddings. The reference pass is modulated at a fixed timestep. + timestep_input = timestep * 0 + 1 if kv_cache_mode == "extract" else timestep + temb = self.time_embedding( + self.timesteps_proj(timestep_input).to(dtype=next(self.time_embedding.parameters()).dtype) + ) + timestep_proj = self.time_projection(temb).unflatten(1, (6, self.dim)) + + encoder_hidden_states = self.text_embedding( + torch.stack( + [torch.cat([u, u.new_zeros(self.text_len - u.size(0), u.size(1))]) for u in encoder_hidden_states] + ) + ) + encoder_hidden_states_image = self.img_emb(encoder_hidden_states_image) if self.use_img_emb else None + + # 4. Per-mode block arguments. + if kv_cache_mode == "extract": + block_kwargs = { + "rotary_emb": self.freqs_ref, + "grid_sizes": grid_sizes, + } + else: + self.freqs = torch.cat( + [ + rope_params(512, d - 4 * (d // 6)), + rope_params(512, 2 * (d // 6)), + rope_params(512, 2 * (d // 6)), + ], + dim=1, + ) + if self.freqs.device != device: + self.freqs = self.freqs.to(device) + + # Latent geometry of the full video, which is what the block mask is built over. The + # current segment is scattered into a buffer of this size inside the attention processor. + origin_latent_frames = origin_len // 4 + 1 + origin_latent_hw = origin_area[0] * origin_area[1] // 256 + + block_mask_id = (origin_latent_frames, origin_latent_hw) + if block_mask_id not in self.block_masks: + self.block_masks[block_mask_id] = self.create_mask( + origin_latent_frames, origin_latent_hw, hidden_states.device + ) + + block_kwargs = { + "rotary_emb": self.freqs, + "grid_sizes": grid_sizes, + "reference_rotary_emb": self.freqs_ref, + "reference_grid_sizes": reference_grid_sizes, + "attention_mask": self.block_masks[block_mask_id], + "origin_latent_frames": origin_latent_frames, + "origin_latent_hw": origin_latent_hw, + } + + block_kwargs.update( + temb=timestep_proj, + encoder_hidden_states=encoder_hidden_states, + encoder_hidden_states_image=encoder_hidden_states_image, + kv_cache_mode=kv_cache_mode, + ) + + # 5. Transformer blocks. + for idx, block in enumerate(self.blocks): + if is_uncondtion and idx == 9: + continue + if torch.is_grad_enabled() and self.gradient_checkpointing: + hidden_states = self._gradient_checkpointing_func( + block, hidden_states, kv_cache=kv_cache.get(idx), **block_kwargs + ) + else: + hidden_states = block(hidden_states, kv_cache=kv_cache.get(idx), **block_kwargs) + + # 6. Output. Under `kv_cache_mode="extract"` the meaningful product is the populated + # cache; the sample is returned anyway so both modes have the same return type. + hidden_states = self.head(hidden_states, temb) + output = [u.float() for u in self.unpatchify(hidden_states, grid_sizes)] + + if not return_dict: + return (output,) + return Transformer2DModelOutput(sample=output) + + def unpatchify(self, x, grid_sizes): + c = self.out_dim + out = [] + for u, v in zip(x, grid_sizes.tolist()): + u = u[: math.prod(v)].view(*v, *self.patch_size, c) + u = torch.einsum("fhwpqrc->cfphqwr", u) + u = u.reshape(c, *[i * j for i, j in zip(v, self.patch_size)]) + out.append(u) + return out diff --git a/src/diffusers/modular_pipelines/__init__.py b/src/diffusers/modular_pipelines/__init__.py index e8d3c71d0fff..0dc5e0b17740 100644 --- a/src/diffusers/modular_pipelines/__init__.py +++ b/src/diffusers/modular_pipelines/__init__.py @@ -47,6 +47,12 @@ ] _import_structure["stable_diffusion_xl"] = ["StableDiffusionXLAutoBlocks", "StableDiffusionXLModularPipeline"] _import_structure["stable_diffusion_3"] = ["StableDiffusion3AutoBlocks", "StableDiffusion3ModularPipeline"] + _import_structure["wan_animate_2"] = [ + "WanAnimate2Blocks", + "WanAnimate2DistilledBlocks", + "WanAnimate2DistilledModularPipeline", + "WanAnimate2ModularPipeline", + ] _import_structure["wan"] = [ "WanBlocks", "Wan22Blocks", @@ -223,6 +229,12 @@ WanImage2VideoModularPipeline, WanModularPipeline, ) + from .wan_animate_2 import ( + WanAnimate2Blocks, + WanAnimate2DistilledBlocks, + WanAnimate2DistilledModularPipeline, + WanAnimate2ModularPipeline, + ) from .z_image import ZImageAutoBlocks, ZImageModularPipeline else: import sys diff --git a/src/diffusers/modular_pipelines/modular_pipeline.py b/src/diffusers/modular_pipelines/modular_pipeline.py index 9880c3fc2502..88cea1b78b1d 100644 --- a/src/diffusers/modular_pipelines/modular_pipeline.py +++ b/src/diffusers/modular_pipelines/modular_pipeline.py @@ -132,6 +132,8 @@ def _helios_pyramid_map_fn(config_dict=None): ("stable-diffusion-xl", _create_default_map_fn("StableDiffusionXLModularPipeline")), ("stable-diffusion-3", _create_default_map_fn("StableDiffusion3ModularPipeline")), ("wan", _wan_map_fn), + ("wan-animate-2", _create_default_map_fn("WanAnimate2ModularPipeline")), + ("wan-animate-2-distilled", _create_default_map_fn("WanAnimate2DistilledModularPipeline")), ("wan-i2v", _wan_i2v_map_fn), ("flux", _create_default_map_fn("FluxModularPipeline")), ("flux-kontext", _create_default_map_fn("FluxKontextModularPipeline")), diff --git a/src/diffusers/modular_pipelines/wan_animate_2/__init__.py b/src/diffusers/modular_pipelines/wan_animate_2/__init__.py new file mode 100644 index 000000000000..d4c196fe09e3 --- /dev/null +++ b/src/diffusers/modular_pipelines/wan_animate_2/__init__.py @@ -0,0 +1,57 @@ +from typing import TYPE_CHECKING + +from ...utils import ( + DIFFUSERS_SLOW_IMPORT, + OptionalDependencyNotAvailable, + _LazyModule, + get_objects_from_module, + is_torch_available, + is_transformers_available, +) + + +_dummy_objects = {} +_import_structure = {} + +try: + if not (is_transformers_available() and is_torch_available()): + raise OptionalDependencyNotAvailable() +except OptionalDependencyNotAvailable: + from ...utils import dummy_torch_and_transformers_objects # noqa F403 + + _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects)) +else: + _import_structure["modular_blocks_wan_animate_2"] = ["WanAnimate2Blocks"] + _import_structure["modular_blocks_wan_animate_2_distilled"] = ["WanAnimate2DistilledBlocks"] + _import_structure["modular_pipeline"] = [ + "WanAnimate2DistilledModularPipeline", + "WanAnimate2ModularPipeline", + ] + _import_structure["video_processor"] = ["WanAnimate2VideoProcessor"] + +if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: + try: + if not (is_transformers_available() and is_torch_available()): + raise OptionalDependencyNotAvailable() + except OptionalDependencyNotAvailable: + from ...utils.dummy_torch_and_transformers_objects import * # noqa F403 + else: + from .modular_blocks_wan_animate_2 import WanAnimate2Blocks + from .modular_blocks_wan_animate_2_distilled import WanAnimate2DistilledBlocks + from .modular_pipeline import ( + WanAnimate2DistilledModularPipeline, + WanAnimate2ModularPipeline, + ) + from .video_processor import WanAnimate2VideoProcessor +else: + import sys + + sys.modules[__name__] = _LazyModule( + __name__, + globals()["__file__"], + _import_structure, + module_spec=__spec__, + ) + + for name, value in _dummy_objects.items(): + setattr(sys.modules[__name__], name, value) diff --git a/src/diffusers/modular_pipelines/wan_animate_2/before_denoise.py b/src/diffusers/modular_pipelines/wan_animate_2/before_denoise.py new file mode 100644 index 000000000000..0ad038e8ccaa --- /dev/null +++ b/src/diffusers/modular_pipelines/wan_animate_2/before_denoise.py @@ -0,0 +1,112 @@ +# 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. + +import math + +import numpy as np +import torch + +from ...utils import logging +from ..modular_pipeline import ModularPipelineBlocks, PipelineState +from ..modular_pipeline_utils import InputParam, OutputParam + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +class WanAnimate2PrepareSegmentsStep(ModularPipelineBlocks): + model_name = "wan-animate-2" + + @property + def description(self) -> str: + return ( + "Step that computes the segment-invariant geometry for the segment loop. The zigzag padding makes " + "every segment exactly `segment_frame_length` frames, so the latent grid, the packed sequence lengths, and the " + "noise shape are the same for all segments and are computed once here." + ) + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + "segment_frame_length", + type_hint=int, + default=81, + description="The number of frames in each inference segment", + ), + InputParam( + "reference_image_latents", + required=True, + type_hint=torch.Tensor, + description="The reference conditioning tensor `[20, 1, latent_height, latent_width]`; " + "provides the latent grid", + ), + InputParam( + "driving_video_pixels", + required=True, + type_hint=torch.Tensor, + description="The preprocessed driving video `[1, 3, T, H, W]`, from the video preprocess step", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + "grid_sizes_ref", + type_hint=torch.Tensor, + description="Post-patch latent grid `[[T, H/2, W/2]]` of a driving-video segment, used as the " + "offset grid of the reference-extraction pass and the reference grid of the denoising passes", + ), + OutputParam( + "max_seq_len", + type_hint=int, + description="Packed sequence length of the generation tokens", + ), + OutputParam( + "max_seq_len_ref", + type_hint=int, + description="Packed sequence length of the reference tokens", + ), + ] + + @torch.no_grad() + def __call__(self, components, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + + latent_height, latent_width = block_state.reference_image_latents.shape[-2:] + + expected = ( + latent_height * components.vae_scale_factor_spatial, + latent_width * components.vae_scale_factor_spatial, + ) + if tuple(block_state.driving_video_pixels.shape[-2:]) != expected: + raise ValueError( + f"`driving_video_pixels` is letterboxed to {tuple(block_state.driving_video_pixels.shape[-2:])} but " + f"the reference image conditioning is {expected} — the video and image preprocess steps must use the " + "same `height`/`width`." + ) + + latent_segment_frames = (block_state.segment_frame_length - 1) // components.vae_scale_factor_temporal + 1 + ref_shape = [latent_segment_frames, latent_height, latent_width] + ref_shape_post = [ref_shape[0], ref_shape[1] // 2, ref_shape[2] // 2] + block_state.grid_sizes_ref = torch.tensor([ref_shape_post], dtype=torch.long) + + # The noise tensor carries one extra latent frame: the reference image's slot. + latent_noise_frames = latent_segment_frames + 1 + block_state.max_seq_len = int(math.ceil(np.prod([latent_noise_frames, latent_height // 2, latent_width // 2]))) + block_state.max_seq_len_ref = int(math.ceil(np.prod(ref_shape) // 4)) + + self.set_block_state(state, block_state) + return components, state diff --git a/src/diffusers/modular_pipelines/wan_animate_2/decoders.py b/src/diffusers/modular_pipelines/wan_animate_2/decoders.py new file mode 100644 index 000000000000..5317ee85e67a --- /dev/null +++ b/src/diffusers/modular_pipelines/wan_animate_2/decoders.py @@ -0,0 +1,99 @@ +# 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. + +import numpy as np +import PIL.Image +import torch + +from ...configuration_utils import FrozenDict +from ...utils import logging +from ..modular_pipeline import ModularPipelineBlocks, PipelineState +from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam +from .video_processor import WanAnimate2VideoProcessor + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +class WanAnimate2DecodeStep(ModularPipelineBlocks): + model_name = "wan-animate-2" + + @property + def description(self) -> str: + return ( + "Step that assembles the final video from the per-segment decoded frames: concatenates the segments, " + "trims the zigzag padding, crops the reference image's letterbox bars back off, and postprocesses. " + "The VAE decode itself happens per segment inside the denoise loop, because each segment conditions " + "on the previous segment's decoded pixels." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec( + "video_processor", + WanAnimate2VideoProcessor, + config=FrozenDict({"vae_scale_factor": 8, "spatial_patch_size": (2, 2), "resample": "bilinear"}), + default_creation_method="from_config", + ), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + "segment_frames", + required=True, + type_hint=list[torch.Tensor], + description="Per-segment decoded frames from the segment denoise loop, each `[1, 3, T, H, W]`", + ), + InputParam( + "real_frame_len", + required=True, + type_hint=int, + description="Number of frames in the driving video before zigzag padding; the output is trimmed to it", + ), + InputParam( + "crop_region", + required=True, + type_hint=tuple[int, int, int, int], + description="`(top, left, height, width)` of the reference image content inside the letterboxed " + "frame, from the image preprocess step", + ), + InputParam( + "output_type", default="np", type_hint=str, description="The output type of the decoded videos" + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + "videos", + type_hint=list[list[PIL.Image.Image]] | list[torch.Tensor] | list[np.ndarray], + description="The generated videos, can be a PIL.Image.Image, torch.Tensor or a numpy array", + ) + ] + + @torch.no_grad() + def __call__(self, components, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + + video = torch.cat(block_state.segment_frames, dim=2)[:, :, : block_state.real_frame_len] + crop_top, crop_left, crop_height, crop_width = block_state.crop_region + video = video[:, :, :, crop_top : crop_top + crop_height, crop_left : crop_left + crop_width] + block_state.videos = components.video_processor.postprocess_video(video, output_type=block_state.output_type) + + self.set_block_state(state, block_state) + return components, state diff --git a/src/diffusers/modular_pipelines/wan_animate_2/denoise.py b/src/diffusers/modular_pipelines/wan_animate_2/denoise.py new file mode 100644 index 000000000000..d96b8f814239 --- /dev/null +++ b/src/diffusers/modular_pipelines/wan_animate_2/denoise.py @@ -0,0 +1,811 @@ +# 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. + +import inspect + +import torch +import torch.nn.functional as F +from tqdm import tqdm + +from ...configuration_utils import FrozenDict +from ...guiders import ClassifierFreeGuidance +from ...models import AutoencoderKLWan, WanAnimate2Transformer3DModel +from ...models.transformers.transformer_wan_animate_2 import WanAnimate2KVCache +from ...schedulers.scheduling_utils import SchedulerMixin +from ...utils import logging +from ...utils.torch_utils import randn_tensor +from ..modular_pipeline import BlockState, LoopSequentialPipelineBlocks, ModularPipelineBlocks, PipelineState +from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam +from .encoders import encode_vae, get_i2v_mask + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +def decode_vae(vae: AutoencoderKLWan, latents: torch.Tensor) -> torch.Tensor: + """De-standardize latents and VAE-decode them to `[B, 3, T, H, W]` pixels in `[-1, 1]`.""" + latents = latents.to(vae.dtype) + latents_mean = ( + torch.tensor(vae.config.latents_mean).view(1, vae.config.z_dim, 1, 1, 1).to(latents.device, latents.dtype) + ) + latents_recip_std = 1.0 / torch.tensor(vae.config.latents_std).view(1, vae.config.z_dim, 1, 1, 1).to( + latents.device, latents.dtype + ) + latents = latents / latents_recip_std + latents_mean + return vae.decode(latents, return_dict=False)[0] + + +# ======================================== +# Segment Loop Leaf Blocks +# ======================================== + + +class WanAnimate2SegmentVaeEncoderStep(ModularPipelineBlocks): + model_name = "wan-animate-2" + + @property + def description(self) -> str: + return ( + "Step within the segment loop that VAE-encodes this segment's slice of the driving video and stacks " + "the i2v conditioning mask on top. The Wan VAE is causal in time, so encoding the whole video once " + "and slicing the latents would not be equivalent — each segment restarts the temporal convolution on " + "its own slice. A streaming mode would replace this block with one fed segments incrementally. This " + "block should be used to compose the `sub_blocks` attribute of `WanAnimate2SegmentLoopWrapper`." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("vae", AutoencoderKLWan), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + "driving_video_pixels", + required=True, + type_hint=torch.Tensor, + description="The preprocessed driving video `[1, 3, T, H, W]`, from the video preprocess step", + ), + InputParam( + "reference_image_latents", + required=True, + type_hint=torch.Tensor, + description="i2v mask + reference image latents `[20, 1, latent_height, latent_width]`, from the image VAE encoder step", + ), + InputParam( + "effective_segment", + required=True, + type_hint=int, + description="Frames each segment advances: `segment_frame_length - prev_segment_conditioning_frames`, from the video preprocess step", + ), + InputParam( + "segment_frame_length", + type_hint=int, + default=81, + description="The number of frames in each inference segment", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + "driving_video_latents", + type_hint=torch.Tensor, + description="VAE latents of this segment's driving-video slice", + ), + OutputParam( + "driving_video_condition", + type_hint=torch.Tensor, + description="i2v mask + driving-slice latents, conditioning the reference-extraction pass", + ), + ] + + @torch.no_grad() + def __call__(self, components, block_state: BlockState, k: int): + device = components._execution_device + + latent_height, latent_width = block_state.reference_image_latents.shape[-2:] + + start = k * block_state.effective_segment + block_state.driving_video_latents = encode_vae( + components.vae, block_state.driving_video_pixels[:, :, start : start + block_state.segment_frame_length] + ) + + condition_mask = get_i2v_mask( + block_state.driving_video_latents.shape[2], + latent_height, + latent_width, + block_state.segment_frame_length, + device=device, + ).to(block_state.driving_video_latents.dtype) + block_state.driving_video_condition = torch.cat([condition_mask, block_state.driving_video_latents[0]], dim=0) + + return components, block_state + + +class WanAnimate2SegmentPrevFramesStep(ModularPipelineBlocks): + model_name = "wan-animate-2" + + @property + def description(self) -> str: + return ( + "Step within the segment loop that builds the generation-side conditioning tensor `reference_latents`: the previous " + "segment's tail frames (zeros for the first segment) are VAE-encoded, masked, and stacked under the " + "reference half `reference_image_latents`. This is how motion continuity crosses segment boundaries — in pixel space, " + "not latent space. This block should be used to compose the `sub_blocks` attribute of " + "`WanAnimate2SegmentLoopWrapper`." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("vae", AutoencoderKLWan), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + "reference_image_latents", + required=True, + type_hint=torch.Tensor, + description="i2v mask + reference image latents `[20, 1, latent_height, latent_width]`, from the image VAE encoder step", + ), + InputParam( + "segment_frame_length", + type_hint=int, + default=81, + description="The number of frames in each inference segment", + ), + InputParam( + "prev_segment_conditioning_frames", + type_hint=int, + default=1, + description="The number of conditioning frames carried over from the previous segment", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + "reference_latents", + type_hint=torch.Tensor, + description="The full conditioning tensor: reference half stacked over the segment half", + ), + ] + + @torch.no_grad() + def __call__(self, components, block_state: BlockState, k: int): + # `block_state.out_frames` is seeded by the loop wrapper and written by the decode step of the + # previous iteration. + device = components._execution_device + + latent_height, latent_width = block_state.reference_image_latents.shape[-2:] + height = latent_height * components.vae_scale_factor_spatial + width = latent_width * components.vae_scale_factor_spatial + + num_frames = block_state.segment_frame_length + 1 + mask_len = block_state.prev_segment_conditioning_frames if k > 0 else 0 + if mask_len > 0: + prev_frames = block_state.out_frames[0, :, -mask_len:].clone().detach() + prev_frames = F.interpolate(prev_frames.permute(1, 0, 2, 3), size=(height, width), mode="bicubic").permute( + 1, 0, 2, 3 + ) + cond_pixels = torch.cat( + [ + prev_frames, + torch.zeros(3, num_frames - mask_len - 1, height, width, device=device), + ], + dim=1, + ) + else: + cond_pixels = torch.zeros(3, num_frames - 1, height, width, device=device) + + prev_segment_cond_latents = encode_vae(components.vae, cond_pixels.unsqueeze(0)).squeeze(0) + prev_segment_cond_mask = get_i2v_mask( + prev_segment_cond_latents.shape[1], latent_height, latent_width, mask_len, device=device + ).to(prev_segment_cond_latents.dtype) + prev_segment_cond_latents = torch.cat([prev_segment_cond_mask, prev_segment_cond_latents], dim=0) + + block_state.reference_latents = torch.cat( + [block_state.reference_image_latents, prev_segment_cond_latents], dim=1 + ) + + return components, block_state + + +class WanAnimate2SegmentPrepareStep(ModularPipelineBlocks): + model_name = "wan-animate-2" + + @property + def description(self) -> str: + return ( + "Step within the segment loop that draws this segment's initial noise and allocates a fresh KV cache " + "for the reference-extraction pass. This block should be used to compose the `sub_blocks` attribute " + "of `WanAnimate2SegmentLoopWrapper`." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("transformer", WanAnimate2Transformer3DModel), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("generator"), + InputParam( + "reference_latents", + required=True, + type_hint=torch.Tensor, + description="The full conditioning tensor: reference half stacked over the segment half", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam("latents", type_hint=torch.Tensor, description="This segment's initial noise"), + OutputParam( + "kv_cache", + type_hint=WanAnimate2KVCache, + description="Fresh per-segment cache for the reference K/V", + ), + ] + + @torch.no_grad() + def __call__(self, components, block_state: BlockState, k: int): + device = components._execution_device + + block_state.latents = randn_tensor( + ( + components.num_channels_latents, + block_state.reference_latents.shape[1], + block_state.reference_latents.shape[-2], + block_state.reference_latents.shape[-1], + ), + generator=block_state.generator, + device=device, + dtype=torch.float32, + ) + block_state.kv_cache = WanAnimate2KVCache(components.transformer.config.num_layers) + + return components, block_state + + +class WanAnimate2SegmentSchedulerResetStep(ModularPipelineBlocks): + model_name = "wan-animate-2" + + @property + def description(self) -> str: + return ( + "Step within the segment loop that resets the scheduler: each segment is an independent denoising " + "trajectory, so the solver state and timesteps are re-prepared per segment. This block should be used " + "to compose the `sub_blocks` attribute of `WanAnimate2SegmentLoopWrapper`." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("scheduler", SchedulerMixin), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("num_inference_steps", default=40), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam("timesteps", type_hint=torch.Tensor, description="This segment's denoising timesteps"), + ] + + @torch.no_grad() + def __call__(self, components, block_state: BlockState, k: int): + device = components._execution_device + + components.scheduler.set_timesteps(block_state.num_inference_steps, device=device) + block_state.timesteps = components.scheduler.timesteps + + return components, block_state + + +class WanAnimate2RefExtractStep(ModularPipelineBlocks): + model_name = "wan-animate-2" + + @property + def description(self) -> str: + return ( + "Step within the segment loop that runs the transformer's reference-extraction pass " + '(`kv_cache_mode="extract"`): the driving-video segment is encoded once and every layer\'s reference ' + "K/V is stored in the KV cache, which the denoising forwards then attend over. This block should be " + "used to compose the `sub_blocks` attribute of `WanAnimate2SegmentLoopWrapper`." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("transformer", WanAnimate2Transformer3DModel), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + "driving_video_latents", + required=True, + type_hint=torch.Tensor, + description="VAE latents of this segment's driving-video slice", + ), + InputParam( + "driving_video_condition", + required=True, + type_hint=torch.Tensor, + description="i2v mask + driving-slice latents, conditioning the reference-extraction pass", + ), + InputParam( + "condition_clip_context", + required=True, + type_hint=torch.Tensor, + description="CLIP vision features of the driving video's first frame", + ), + InputParam( + "prompt_ref_embeds", + required=True, + type_hint=torch.Tensor, + description="Text embeddings of the reference prompt, guiding the reference-extraction pass", + ), + InputParam( + "kv_cache", + required=True, + type_hint=WanAnimate2KVCache, + description="Per-segment cache holding every layer's reference K/V", + ), + InputParam( + "timesteps", + required=True, + type_hint=torch.Tensor, + description="This segment's denoising timesteps", + ), + InputParam( + "max_seq_len_ref", + required=True, + type_hint=int, + description="Packed sequence length of the reference tokens", + ), + InputParam( + "grid_sizes_ref", + required=True, + type_hint=torch.Tensor, + description="Post-patch latent grid `[[T, H/2, W/2]]` of a driving-video segment", + ), + ] + + @torch.no_grad() + def __call__(self, components, block_state: BlockState, k: int): + device = components._execution_device + transformer_dtype = components.transformer.dtype + + t_ref = torch.tensor([block_state.timesteps[0].item()], device=device, dtype=transformer_dtype) + components.transformer( + [block_state.driving_video_latents[0].to(transformer_dtype)], + timestep=t_ref, + encoder_hidden_states=[block_state.prompt_ref_embeds[0].to(transformer_dtype)], + encoder_hidden_states_image=block_state.condition_clip_context.to(transformer_dtype), + condition_latents=[block_state.driving_video_condition.to(transformer_dtype)], + kv_cache=block_state.kv_cache, + kv_cache_mode="extract", + seq_len=block_state.max_seq_len_ref, + offset_grid_sizes=block_state.grid_sizes_ref, + ) + + return components, block_state + + +# ======================================== +# Inner Denoising Blocks +# ======================================== + + +class WanAnimate2SegmentDenoiseInner(ModularPipelineBlocks): + model_name = "wan-animate-2" + + @property + def description(self) -> str: + return ( + "Inner timestep loop that denoises one segment with guidance, attending over the segment's cached " + "reference K/V. The unconditional branch passes `is_uncondtion=True` to the transformer (it skips a " + "dedicated layer on that branch), routed through the guider as a per-branch input. This block should " + "be used to compose the `sub_blocks` attribute of `WanAnimate2SegmentLoopWrapper`." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("transformer", WanAnimate2Transformer3DModel), + ComponentSpec("scheduler", SchedulerMixin), + ComponentSpec( + "guider", + ClassifierFreeGuidance, + config=FrozenDict({"guidance_scale": 3.0}), + default_creation_method="from_config", + ), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + "latents", + required=True, + type_hint=torch.Tensor, + description="This segment's latents", + ), + InputParam( + "reference_latents", + required=True, + type_hint=torch.Tensor, + description="The full conditioning tensor: reference half stacked over the segment half", + ), + InputParam( + "kv_cache", + required=True, + type_hint=WanAnimate2KVCache, + description="Per-segment cache holding every layer's reference K/V", + ), + InputParam( + "timesteps", + required=True, + type_hint=torch.Tensor, + description="This segment's denoising timesteps", + ), + InputParam.template("num_inference_steps", default=40), + InputParam( + "num_segments", + required=True, + type_hint=int, + description="Total number of segments in the driving video, from the video preprocess step", + ), + InputParam( + "max_seq_len", + required=True, + type_hint=int, + description="Packed sequence length of the generation tokens", + ), + InputParam( + "grid_sizes_ref", + required=True, + type_hint=torch.Tensor, + description="Post-patch latent grid `[[T, H/2, W/2]]` of a driving-video segment", + ), + InputParam( + "segment_frame_length", + type_hint=int, + default=81, + description="The number of frames in each inference segment", + ), + InputParam( + "height", + required=True, + type_hint=int, + description="The resolved frame height in pixels", + ), + InputParam( + "width", + required=True, + type_hint=int, + description="The resolved frame width in pixels", + ), + InputParam.template("generator"), + InputParam.template("prompt_embeds"), + InputParam.template("negative_prompt_embeds"), + InputParam.template("denoiser_input_fields"), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam.template("latents"), + ] + + @torch.no_grad() + def __call__(self, components, block_state: BlockState, k: int): + transformer_dtype = components.transformer.dtype + + guider_inputs = { + "encoder_hidden_states": (block_state.prompt_embeds, block_state.negative_prompt_embeds), + "is_uncondtion": (False, True), + } + + # Everything the transformer accepts from the tagged conditioning fields, minus what the guider manages + # per branch (currently the reference image's CLIP features, `encoder_hidden_states_image`). + transformer_args = set(inspect.signature(components.transformer.forward).parameters.keys()) + shared_kwargs = { + name: value.to(transformer_dtype) if isinstance(value, torch.Tensor) else value + for name, value in block_state.denoiser_input_fields.items() + if name in transformer_args and name not in guider_inputs + } + + with tqdm( + total=len(block_state.timesteps), desc=f"Segment {k + 1}/{block_state.num_segments}" + ) as progress_bar: + for i, t in enumerate(block_state.timesteps): + timestep = torch.stack([t]) + + components.guider.set_state(step=i, num_inference_steps=block_state.num_inference_steps, timestep=t) + guider_state = components.guider.prepare_inputs(guider_inputs) + + for guider_state_batch in guider_state: + components.guider.prepare_models(components.transformer) + + guider_state_batch.noise_pred = components.transformer( + [block_state.latents.to(transformer_dtype)], + timestep=timestep, + encoder_hidden_states=[guider_state_batch.encoder_hidden_states[0].to(transformer_dtype)], + condition_latents=[block_state.reference_latents.to(transformer_dtype)], + kv_cache=block_state.kv_cache, + kv_cache_mode="cached", + seq_len=block_state.max_seq_len, + reference_grid_sizes=block_state.grid_sizes_ref, + origin_len=block_state.segment_frame_length, + origin_area=[block_state.height, block_state.width], + is_uncondtion=guider_state_batch.is_uncondtion, + **shared_kwargs, + ).sample[0] + + components.guider.cleanup_models(components.transformer) + + noise_pred = components.guider(guider_state)[0] + + latents = components.scheduler.step( + noise_pred.unsqueeze(0), + t, + block_state.latents.unsqueeze(0), + return_dict=False, + generator=block_state.generator, + )[0] + block_state.latents = latents.squeeze(0) + + progress_bar.update() + + return components, block_state + + +class WanAnimate2DistilledSegmentDenoiseInner(WanAnimate2SegmentDenoiseInner): + model_name = "wan-animate-2-distilled" + + @property + def description(self) -> str: + return ( + "Inner timestep loop that denoises one segment for the distilled model, which is trained for few-step " + "sampling without classifier-free guidance — the guider defaults to `guidance_scale=1.0`, so only the " + "conditional branch runs. This block should be used to compose the `sub_blocks` attribute of " + "`WanAnimate2SegmentLoopWrapper`." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("transformer", WanAnimate2Transformer3DModel), + ComponentSpec("scheduler", SchedulerMixin), + ComponentSpec( + "guider", + ClassifierFreeGuidance, + config=FrozenDict({"guidance_scale": 1.0}), + default_creation_method="from_config", + ), + ] + + +# ======================================== +# Post-Denoise +# ======================================== + + +class WanAnimate2SegmentDecodeStep(ModularPipelineBlocks): + model_name = "wan-animate-2" + + @property + def description(self) -> str: + return ( + "Step within the segment loop that VAE-decodes the denoised segment. Decoding happens inside the loop " + "because the next segment conditions on this segment's decoded pixels. Finished frames move to CPU and " + "the per-segment KV cache and latents are freed — holding them across segments fragments the " + "allocator enough to OOM at high resolution. This block should be used to compose the `sub_blocks` " + "attribute of `WanAnimate2SegmentLoopWrapper`." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("vae", AutoencoderKLWan), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + "latents", + required=True, + type_hint=torch.Tensor, + description="This segment's latents", + ), + InputParam( + "kv_cache", + required=True, + type_hint=WanAnimate2KVCache, + description="Per-segment cache holding every layer's reference K/V", + ), + InputParam( + "prev_segment_conditioning_frames", + type_hint=int, + default=1, + description="The number of conditioning frames carried over from the previous segment", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + "out_frames", + type_hint=torch.Tensor, + description="This segment's decoded frames on device; the next segment conditions on its tail", + ), + ] + + @torch.no_grad() + def __call__(self, components, block_state: BlockState, k: int): + latents = block_state.latents.to(torch.float32) + # The first latent frame is the reference image's slot, not video content. + out_frames = decode_vae(components.vae, latents[:, 1:]) + + if k > 0: + out_frames = out_frames[:, :, block_state.prev_segment_conditioning_frames :] + + block_state.segment_frames.append(out_frames.cpu()) + block_state.out_frames = out_frames + + block_state.kv_cache.clear() + block_state.kv_cache = None + block_state.latents = None + torch.cuda.empty_cache() + + return components, block_state + + +# ======================================== +# Segment Loop Wrapper +# ======================================== + + +class WanAnimate2SegmentLoopWrapper(LoopSequentialPipelineBlocks): + model_name = "wan-animate-2" + + @property + def description(self) -> str: + return ( + "Pipeline block that iterates over the driving video's segments. At each segment it runs sub-blocks " + "for per-segment encoding, preparation, reference extraction, denoising, and decoding; each segment " + "conditions on the previous one's decoded tail frames." + ) + + @property + def loop_inputs(self) -> list[InputParam]: + return [ + InputParam( + "num_segments", + required=True, + type_hint=int, + description="Total number of segments in the driving video, from the video preprocess step", + ), + ] + + @property + def loop_intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + "segment_frames", + type_hint=list[torch.Tensor], + description="Per-segment decoded frames on CPU, each `[1, 3, T, H, W]`", + ), + ] + + @torch.no_grad() + def __call__(self, components, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + + # Seed the loop-carried state: `segment_frames` collects each segment's decoded frames (the decode step + # appends to it); `out_frames` is the previous segment's decoded frames — written by the decode step, read + # by the prev-frames step of the next iteration. `None` marks "no previous segment" for the first iteration. + block_state.segment_frames = [] + block_state.out_frames = None + + for k in range(block_state.num_segments): + components, block_state = self.loop_step(components, block_state, k=k) + + self.set_block_state(state, block_state) + return components, state + + +# ======================================== +# Composed Segment Denoise Steps +# ======================================== + + +class WanAnimate2DenoiseStep(WanAnimate2SegmentLoopWrapper): + block_classes = [ + WanAnimate2SegmentVaeEncoderStep, + WanAnimate2SegmentPrevFramesStep, + WanAnimate2SegmentPrepareStep, + WanAnimate2SegmentSchedulerResetStep, + WanAnimate2RefExtractStep, + WanAnimate2SegmentDenoiseInner, + WanAnimate2SegmentDecodeStep, + ] + block_names = [ + "vae_encoder", + "prev_frames", + "prepare", + "scheduler_reset", + "ref_extract", + "denoise_inner", + "decode", + ] + + @property + def description(self) -> str: + return ( + "Segment denoise step that iterates over the driving video's segments.\n" + "At each segment: vae_encoder -> prev_frames -> prepare -> scheduler_reset -> ref_extract -> " + "denoise_inner -> decode." + ) + + +class WanAnimate2DistilledDenoiseStep(WanAnimate2SegmentLoopWrapper): + model_name = "wan-animate-2-distilled" + + block_classes = [ + WanAnimate2SegmentVaeEncoderStep, + WanAnimate2SegmentPrevFramesStep, + WanAnimate2SegmentPrepareStep, + WanAnimate2SegmentSchedulerResetStep, + WanAnimate2RefExtractStep, + WanAnimate2DistilledSegmentDenoiseInner, + WanAnimate2SegmentDecodeStep, + ] + block_names = [ + "vae_encoder", + "prev_frames", + "prepare", + "scheduler_reset", + "ref_extract", + "denoise_inner", + "decode", + ] + + @property + def description(self) -> str: + return ( + "Segment denoise step for the distilled model that iterates over the driving video's segments.\n" + "At each segment: vae_encoder -> prev_frames -> prepare -> scheduler_reset -> ref_extract -> " + "denoise_inner (no classifier-free guidance) -> decode." + ) diff --git a/src/diffusers/modular_pipelines/wan_animate_2/encoders.py b/src/diffusers/modular_pipelines/wan_animate_2/encoders.py new file mode 100644 index 000000000000..21b70f636f7d --- /dev/null +++ b/src/diffusers/modular_pipelines/wan_animate_2/encoders.py @@ -0,0 +1,596 @@ +# 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. + +import math + +import numpy as np +import PIL.Image +import torch +import torch.nn.functional as F +from transformers import AutoTokenizer, CLIPVisionModel, UMT5EncoderModel + +from ...configuration_utils import FrozenDict +from ...models import AutoencoderKLWan +from ...utils import logging +from ..modular_pipeline import ModularPipelineBlocks, PipelineState +from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam +from .video_processor import WanAnimate2VideoProcessor + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +CLIP_MEAN = [0.48145466, 0.4578275, 0.40821073] +CLIP_STD = [0.26862954, 0.26130258, 0.27577711] + + +# Like diffusers.modular_pipelines.wan.encoders.get_t5_prompt_embeds, but without the whitespace +# cleaning -- the reference implementation encodes the prompt exactly as given. +def get_t5_prompt_embeds( + text_encoder: UMT5EncoderModel, + tokenizer: AutoTokenizer, + prompt: str | list[str], + max_sequence_length: int, + device: torch.device, +): + dtype = text_encoder.dtype + prompt = [prompt] if isinstance(prompt, str) else prompt + + text_inputs = tokenizer( + prompt, + padding="max_length", + max_length=max_sequence_length, + truncation=True, + add_special_tokens=True, + return_attention_mask=True, + return_tensors="pt", + ) + text_input_ids, mask = text_inputs.input_ids, text_inputs.attention_mask + seq_lens = mask.gt(0).sum(dim=1).long() + prompt_embeds = text_encoder(text_input_ids.to(device), mask.to(device)).last_hidden_state + prompt_embeds = prompt_embeds.to(dtype=dtype, device=device) + prompt_embeds = [u[:v] for u, v in zip(prompt_embeds, seq_lens)] + prompt_embeds = torch.stack( + [torch.cat([u, u.new_zeros(max_sequence_length - u.size(0), u.size(1))]) for u in prompt_embeds], dim=0 + ) + + return prompt_embeds + + +def clip_visual_encode(image_encoder, tensor, device, dtype): + """Encode tensor to CLIP features (bicubic to 224×224, matching original).""" + if tensor.ndim == 3: + tensor = tensor.unsqueeze(1) + videos = F.interpolate(tensor.transpose(0, 1), size=(224, 224), mode="bicubic", align_corners=False) + videos = videos.mul_(0.5).add_(0.5) + mean = torch.tensor(CLIP_MEAN, device=device, dtype=videos.dtype).view(1, 3, 1, 1) + std = torch.tensor(CLIP_STD, device=device, dtype=videos.dtype).view(1, 3, 1, 1) + videos = (videos - mean) / std + out = image_encoder(pixel_values=videos.to(dtype), output_hidden_states=True) + return out.hidden_states[-2] + + +def get_i2v_mask(lat_t, lat_h, lat_w, mask_len=1, device="cuda"): + """Create an i2v mask in latent space. + + mask_len is in PIXEL space. Returns [4, lat_t, lat_h, lat_w] (no batch dim). + """ + msk = torch.zeros(1, (lat_t - 1) * 4 + 1, lat_h, lat_w, device=device) + msk[:, :mask_len] = 1 + msk = torch.concat([torch.repeat_interleave(msk[:, 0:1], repeats=4, dim=1), msk[:, 1:]], dim=1) + msk = msk.view(1, msk.shape[1] // 4, 4, lat_h, lat_w) + msk = msk.transpose(1, 2)[0] + return msk + + +def get_frame_indices(num_frames, video_fps, target_fps): + """Nearest-neighbour resample of a `video_fps` clip to `target_fps`.""" + num_target_frames = int(num_frames / video_fps * target_fps) + times = np.arange(0, num_target_frames) / target_fps + frame_indices = np.round(times * video_fps).astype(int) + return np.clip(frame_indices, 0, num_frames - 1).tolist() + + +def encode_vae(vae: AutoencoderKLWan, video: torch.Tensor) -> torch.Tensor: + """VAE-encode a `[B, C, T, H, W]` clip (mode of the distribution) and standardize the latents.""" + latents = vae.encode(video.to(vae.dtype)).latent_dist.mode() + latents_mean = ( + torch.tensor(vae.config.latents_mean).view(1, vae.config.z_dim, 1, 1, 1).to(latents.device, latents.dtype) + ) + latents_recip_std = 1.0 / torch.tensor(vae.config.latents_std).view(1, vae.config.z_dim, 1, 1, 1).to( + latents.device, latents.dtype + ) + return (latents - latents_mean) * latents_recip_std + + +# ======================================== +# Text Encoder +# ======================================== + + +class WanAnimate2TextEncoderStep(ModularPipelineBlocks): + model_name = "wan-animate-2" + + @property + def description(self) -> str: + return ( + "Text Encoder step that encodes the character/background prompt, the negative prompt (when the " + "pipeline's guider needs unconditional embeddings, or one is passed explicitly), and the fixed " + "reference prompt for the driving-video context" + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("text_encoder", UMT5EncoderModel), + ComponentSpec("tokenizer", AutoTokenizer), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("prompt"), + InputParam.template("negative_prompt"), + InputParam( + "prompt_ref", + default="人物动作的参考视频", + type_hint=str, + description="The reference prompt for the driving video context", + ), + InputParam.template("max_sequence_length"), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam.template("prompt_embeds"), + OutputParam.template("negative_prompt_embeds"), + OutputParam( + "prompt_ref_embeds", + type_hint=torch.Tensor, + description="text embeddings of the reference prompt, conditioning the reference-extraction pass", + ), + ] + + @staticmethod + def check_inputs(block_state): + if not isinstance(block_state.prompt, str): + raise ValueError(f"`prompt` has to be of type `str` but is {type(block_state.prompt)}") + + @torch.no_grad() + def __call__(self, components, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + self.check_inputs(block_state) + + device = components._execution_device + + block_state.prompt_embeds = get_t5_prompt_embeds( + components.text_encoder, + components.tokenizer, + block_state.prompt, + block_state.max_sequence_length, + device, + ) + # The guider is not a component of this block: when the step runs inside the full pipeline, + # the denoise step's guider determines (via `requires_unconditional_embeds`) whether + # unconditional embeddings are needed, defaulting the negative prompt to "". Standalone, + # there is no guider and the negative prompt is only encoded when the caller passes one. + block_state.negative_prompt_embeds = None + if components.requires_unconditional_embeds or block_state.negative_prompt is not None: + block_state.negative_prompt_embeds = get_t5_prompt_embeds( + components.text_encoder, + components.tokenizer, + block_state.negative_prompt or "", + block_state.max_sequence_length, + device, + ) + block_state.prompt_ref_embeds = get_t5_prompt_embeds( + components.text_encoder, + components.tokenizer, + block_state.prompt_ref, + block_state.max_sequence_length, + device, + ) + + self.set_block_state(state, block_state) + return components, state + + +# ======================================== +# Preprocessing +# ======================================== + + +class WanAnimate2ProcessImagesInputStep(ModularPipelineBlocks): + model_name = "wan-animate-2" + + @property + def description(self) -> str: + return ( + "Image Resize step that resolves the output frame from the target area (`height * width`) and the " + "reference image's aspect ratio, then letterboxes the reference image into that frame. The recorded " + "crop box is used at the end to crop the letterbox bars back off the generated video." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec( + "image_processor", + WanAnimate2VideoProcessor, + config=FrozenDict({"vae_scale_factor": 8, "spatial_patch_size": (2, 2), "resample": "bicubic"}), + default_creation_method="from_config", + ), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("image", description="The reference image holding the character to animate."), + InputParam( + "height", + type_hint=int, + default=800, + description="Together with `width`, the target *area* of the generated video; the aspect ratio " + "comes from `image`. Overwritten with the resolved frame height.", + ), + InputParam( + "width", + type_hint=int, + default=640, + description="See `height`. Overwritten with the resolved frame width.", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + "image_pixels", + type_hint=torch.Tensor, + description="The letterboxed reference image as a `[1, 3, H, W]` tensor in `[-1, 1]`", + ), + OutputParam( + "crop_region", + type_hint=tuple[int, int, int, int], + description="`(top, left, height, width)` of the reference image content inside the letterboxed frame", + ), + ] + + @torch.no_grad() + def __call__(self, components, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + + device = components._execution_device + + image_height, image_width = components.image_processor.get_default_height_width(block_state.image) + mod_value = components.vae_scale_factor_spatial * 2 + aspect_ratio = image_height / image_width + max_area = block_state.height * block_state.width + block_state.height = int(math.sqrt(max_area * aspect_ratio)) // mod_value * mod_value + block_state.width = int(math.sqrt(max_area / aspect_ratio)) // mod_value * mod_value + + height, width = block_state.height, block_state.width + crop_width = width if width / height < image_width / image_height else image_width * height // image_height + crop_height = height if width / height >= image_width / image_height else image_height * width // image_width + crop_top = (height - crop_height) // 2 + crop_left = (width - crop_width) // 2 + block_state.crop_region = (crop_top, crop_left, crop_height, crop_width) + + block_state.image_pixels = components.image_processor.preprocess( + block_state.image, height=height, width=width, resize_mode="fill" + ).to(device, dtype=torch.float32) + + self.set_block_state(state, block_state) + return components, state + + +class WanAnimate2ProcessVideosInputStep(ModularPipelineBlocks): + model_name = "wan-animate-2" + + @property + def description(self) -> str: + return ( + "Video preprocess step that optionally resamples the driving video to the model's frame rate, " + "letterboxes every frame into the resolved output frame, and zigzag-pads the tail so the frame count " + "splits into whole segments. The mirrored padding frames are real content to the model; the surplus " + "generated frames are trimmed off again at the end." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec( + "video_processor", + WanAnimate2VideoProcessor, + config=FrozenDict({"vae_scale_factor": 8, "spatial_patch_size": (2, 2), "resample": "bilinear"}), + default_creation_method="from_config", + ), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + "driving_video", + required=True, + type_hint=list[PIL.Image.Image], + description="The driving video that provides the motion, in any format accepted by " + "`VideoProcessor.preprocess_video`.", + ), + InputParam( + "driving_video_fps", + type_hint=float, + description="The frame rate `driving_video` was captured at — `load_video(..., return_fps=True)` " + "reports it. When set, the driving frames are resampled from it to `fps`; when `None` they are " + "used as-is.", + ), + InputParam("fps", type_hint=int, default=24, description="The frame rate the model generates at"), + InputParam( + "segment_frame_length", + type_hint=int, + default=81, + description="The number of frames in each inference segment", + ), + InputParam( + "prev_segment_conditioning_frames", + type_hint=int, + default=1, + description="The number of conditioning frames carried over from the previous segment", + ), + InputParam( + "height", + type_hint=int, + default=800, + description="The height the driving frames are letterboxed to; must match the reference image's " + "resolved height. In the assembled pipeline the image preprocess step supplies the resolved value.", + ), + InputParam( + "width", + type_hint=int, + default=640, + description="See `height`.", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + "driving_video_pixels", + type_hint=torch.Tensor, + description="The resampled, letterboxed, and zigzag-padded driving video, `[1, 3, T, height, width]` " + "in `[-1, 1]`", + ), + OutputParam( + "real_frame_len", + type_hint=int, + description="Number of driving frames before zigzag padding; the output is trimmed to it", + ), + OutputParam("num_segments", type_hint=int, description="Number of inference segments"), + OutputParam( + "effective_segment", + type_hint=int, + description="Frames each segment advances by (`segment_frame_length - prev_segment_conditioning_frames`)", + ), + ] + + @torch.no_grad() + def __call__(self, components, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + + device = components._execution_device + + # Resample the driving video to the model's frame rate + driving_video = block_state.driving_video + if block_state.driving_video_fps is not None: + frame_indices = get_frame_indices(len(driving_video), block_state.driving_video_fps, block_state.fps) + driving_video = [driving_video[i] for i in frame_indices] + + # each frame letterboxed into the target frame -> [1, 3, T, height, width]` + driving_video = components.video_processor.preprocess_video( + driving_video, height=block_state.height, width=block_state.width, resize_mode="fill" + ).to(device, dtype=torch.float32) + + # Segments overlap by `prev_segment_conditioning_frames`, so each segment advances by + # `effective_segment` new frames. + real_frame_len = driving_video.shape[2] + effective_segment = block_state.segment_frame_length - block_state.prev_segment_conditioning_frames + # If the leftover frames don't fill a whole final segment, pad it with a zigzag pattern: + # frames [0 1 2 3 4] with 3 padding frames -> [0 1 2 3 4 | 4 3 2]. The frames generated + # for the padding are trimmed off again in the decode step. + if real_frame_len > block_state.prev_segment_conditioning_frames: + leftover_frames = (real_frame_len - block_state.prev_segment_conditioning_frames) % effective_segment + else: + leftover_frames = 0 + num_padding = effective_segment - leftover_frames if leftover_frames > 0 else 0 + target_num_frames = real_frame_len + num_padding + + if num_padding > 0: + padding_frames = driving_video[:, :, real_frame_len - num_padding : real_frame_len].flip(2) + driving_video = torch.cat([driving_video, padding_frames], dim=2) + + block_state.driving_video_pixels = driving_video + block_state.real_frame_len = real_frame_len + block_state.effective_segment = effective_segment + block_state.num_segments = ( + target_num_frames - block_state.prev_segment_conditioning_frames + effective_segment - 1 + ) // effective_segment + + self.set_block_state(state, block_state) + return components, state + + +# ======================================== +# Image Encoders (CLIP) +# ======================================== + + +class WanAnimate2ImageClipEncoderStep(ModularPipelineBlocks): + model_name = "wan-animate-2" + + @property + def description(self) -> str: + return "Image Encoder step that computes CLIP vision features of the letterboxed reference image" + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("image_encoder", CLIPVisionModel), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + "image_pixels", + required=True, + type_hint=torch.Tensor, + description="The letterboxed reference image `[1, 3, H, W]` in `[-1, 1]`, from the image preprocess step", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + "encoder_hidden_states_image", + type_hint=torch.Tensor, + kwargs_type="denoiser_input_fields", + description="CLIP vision features of the reference image, conditioning every denoising forward", + ), + ] + + @torch.no_grad() + def __call__(self, components, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + + device = components._execution_device + block_state.encoder_hidden_states_image = clip_visual_encode( + components.image_encoder, block_state.image_pixels[0], device, components.image_encoder.dtype + ) + + self.set_block_state(state, block_state) + return components, state + + +class WanAnimate2VideoClipEncoderStep(ModularPipelineBlocks): + model_name = "wan-animate-2" + + @property + def description(self) -> str: + return ( + "Image Encoder step that computes CLIP vision features of the driving video's first frame, " + "conditioning the per-segment reference-extraction pass" + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("image_encoder", CLIPVisionModel), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + "driving_video_pixels", + required=True, + type_hint=torch.Tensor, + description="The preprocessed driving video `[1, 3, T, H, W]`, from the video preprocess step", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + "condition_clip_context", + type_hint=torch.Tensor, + description="CLIP vision features of the driving video's first frame", + ), + ] + + @torch.no_grad() + def __call__(self, components, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + + device = components._execution_device + block_state.condition_clip_context = clip_visual_encode( + components.image_encoder, block_state.driving_video_pixels[0, :, 0], device, components.image_encoder.dtype + ) + + self.set_block_state(state, block_state) + return components, state + + +# ======================================== +# VAE Encoders +# ======================================== + + +class WanAnimate2ImageVaeEncoderStep(ModularPipelineBlocks): + model_name = "wan-animate-2" + + @property + def description(self) -> str: + return ( + "VAE Encoder step that encodes the letterboxed reference image and stacks the i2v conditioning mask " + "on top, producing the reference half of the conditioning tensor `reference_latents`" + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("vae", AutoencoderKLWan), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + "image_pixels", + required=True, + type_hint=torch.Tensor, + description="The letterboxed reference image `[1, 3, H, W]` in `[-1, 1]`, from the image preprocess step", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + "reference_image_latents", + type_hint=torch.Tensor, + description="i2v mask + reference image latents, `[20, 1, latent_height, latent_width]`", + ), + ] + + @torch.no_grad() + def __call__(self, components, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + + device = components._execution_device + + ref_latents = encode_vae(components.vae, block_state.image_pixels.unsqueeze(2)) + + height, width = block_state.image_pixels.shape[-2:] + latent_height = height // components.vae_scale_factor_spatial + latent_width = width // components.vae_scale_factor_spatial + + mask_ref = get_i2v_mask(1, latent_height, latent_width, 1, device=device).to(ref_latents.dtype) + block_state.reference_image_latents = torch.cat([mask_ref, ref_latents[0]], dim=0) + + self.set_block_state(state, block_state) + return components, state diff --git a/src/diffusers/modular_pipelines/wan_animate_2/modular_blocks_wan_animate_2.py b/src/diffusers/modular_pipelines/wan_animate_2/modular_blocks_wan_animate_2.py new file mode 100644 index 000000000000..f77eb378c15b --- /dev/null +++ b/src/diffusers/modular_pipelines/wan_animate_2/modular_blocks_wan_animate_2.py @@ -0,0 +1,316 @@ +# 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. + +import torch + +from ...utils import logging +from ..modular_pipeline import SequentialPipelineBlocks +from ..modular_pipeline_utils import InsertableDict, OutputParam +from .before_denoise import WanAnimate2PrepareSegmentsStep +from .decoders import WanAnimate2DecodeStep +from .denoise import WanAnimate2DenoiseStep +from .encoders import ( + WanAnimate2ImageClipEncoderStep, + WanAnimate2ImageVaeEncoderStep, + WanAnimate2ProcessImagesInputStep, + WanAnimate2ProcessVideosInputStep, + WanAnimate2TextEncoderStep, + WanAnimate2VideoClipEncoderStep, +) + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +# ==================== +# 1. Encoder groups +# ==================== + + +WanAnimate2ImageEncoderBlocks = InsertableDict( + [ + ("preprocess", WanAnimate2ProcessImagesInputStep()), + ("encode", WanAnimate2ImageClipEncoderStep()), + ] +) + + +# auto_docstring +class WanAnimate2ImageEncodeStep(SequentialPipelineBlocks): + """ + Image encoder step that letterboxes the reference character image to the resolved resolution and CLIP-encodes it + into `encoder_hidden_states_image`. + + Components: + image_processor (`WanAnimate2VideoProcessor`) image_encoder (`CLIPVisionModel`) + + Inputs: + image (`Image | list`): + The reference image holding the character to animate. + height (`int`, *optional*, defaults to 800): + Together with `width`, the target *area* of the generated video; the aspect ratio comes from `image`. + Overwritten with the resolved frame height. + width (`int`, *optional*, defaults to 640): + See `height`. Overwritten with the resolved frame width. + + Outputs: + image_pixels (`Tensor`): + The letterboxed reference image as a `[1, 3, H, W]` tensor in `[-1, 1]` + crop_region (`tuple`): + `(top, left, height, width)` of the reference image content inside the letterboxed frame + encoder_hidden_states_image (`Tensor`): + CLIP vision features of the reference image, conditioning every denoising forward + """ + + model_name = "wan-animate-2" + block_classes = WanAnimate2ImageEncoderBlocks.values() + block_names = WanAnimate2ImageEncoderBlocks.keys() + + @property + def description(self): + return ( + "Image encoder step that letterboxes the reference character image to the resolved resolution and " + "CLIP-encodes it into `encoder_hidden_states_image`." + ) + + +WanAnimate2VideoEncoderBlocks = InsertableDict( + [ + ("preprocess", WanAnimate2ProcessVideosInputStep()), + ("encode", WanAnimate2VideoClipEncoderStep()), + ] +) + + +# auto_docstring +class WanAnimate2VideoEncodeStep(SequentialPipelineBlocks): + """ + Video encoder step that preprocesses the driving video (fps resample, letterbox, zigzag padding to a whole number + of segments) and CLIP-encodes its first frame into `condition_clip_context`. + + Components: + video_processor (`WanAnimate2VideoProcessor`) image_encoder (`CLIPVisionModel`) + + Inputs: + driving_video (`list`): + The driving video that provides the motion, in any format accepted by `VideoProcessor.preprocess_video`. + driving_video_fps (`float`, *optional*): + The frame rate `driving_video` was captured at — `load_video(..., return_fps=True)` reports it. When set, + the driving frames are resampled from it to `fps`; when `None` they are used as-is. + fps (`int`, *optional*, defaults to 24): + The frame rate the model generates at + segment_frame_length (`int`, *optional*, defaults to 81): + The number of frames in each inference segment + prev_segment_conditioning_frames (`int`, *optional*, defaults to 1): + The number of conditioning frames carried over from the previous segment + height (`int`, *optional*, defaults to 800): + The height the driving frames are letterboxed to; must match the reference image's resolved height. In + the assembled pipeline the image preprocess step supplies the resolved value. + width (`int`, *optional*, defaults to 640): + See `height`. + + Outputs: + driving_video_pixels (`Tensor`): + The resampled, letterboxed, and zigzag-padded driving video, `[1, 3, T, height, width]` in `[-1, 1]` + real_frame_len (`int`): + Number of driving frames before zigzag padding; the output is trimmed to it + num_segments (`int`): + Number of inference segments + effective_segment (`int`): + Frames each segment advances by (`segment_frame_length - prev_segment_conditioning_frames`) + condition_clip_context (`Tensor`): + CLIP vision features of the driving video's first frame + """ + + model_name = "wan-animate-2" + block_classes = WanAnimate2VideoEncoderBlocks.values() + block_names = WanAnimate2VideoEncoderBlocks.keys() + + @property + def description(self): + return ( + "Video encoder step that preprocesses the driving video (fps resample, letterbox, zigzag padding to a " + "whole number of segments) and CLIP-encodes its first frame into `condition_clip_context`." + ) + + +# ==================== +# 2. Core denoise +# ==================== + + +WanAnimate2CoreDenoiseBlocks = InsertableDict( + [ + ("prepare_segments", WanAnimate2PrepareSegmentsStep()), + ("denoise", WanAnimate2DenoiseStep()), + ] +) + + +# auto_docstring +class WanAnimate2CoreDenoiseStep(SequentialPipelineBlocks): + """ + Core denoise step that computes the segment-invariant geometry and runs the segment-by-segment denoising loop, + decoding each segment inside the loop because the next segment conditions on its decoded pixels. + + Components: + vae (`AutoencoderKLWan`) transformer (`WanAnimate2Transformer3DModel`) scheduler (`SchedulerMixin`) guider + (`ClassifierFreeGuidance`) + + Inputs: + segment_frame_length (`int`, *optional*, defaults to 81): + The number of frames in each inference segment + reference_image_latents (`Tensor`): + The reference conditioning tensor `[20, 1, latent_height, latent_width]`; provides the latent grid + driving_video_pixels (`Tensor`): + The preprocessed driving video `[1, 3, T, H, W]`, from the video preprocess step + num_segments (`int`): + Total number of segments in the driving video, from the video preprocess step + effective_segment (`int`): + Frames each segment advances: `segment_frame_length - prev_segment_conditioning_frames`, from the video + preprocess step + prev_segment_conditioning_frames (`int`, *optional*, defaults to 1): + The number of conditioning frames carried over from the previous segment + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + num_inference_steps (`int`, *optional*, defaults to 40): + The number of denoising steps. + condition_clip_context (`Tensor`): + CLIP vision features of the driving video's first frame + prompt_ref_embeds (`Tensor`): + Text embeddings of the reference prompt, guiding the reference-extraction pass + height (`int`): + The resolved frame height in pixels + width (`int`): + The resolved frame width in pixels + prompt_embeds (`Tensor`): + text embeddings used to guide the image generation. Can be generated from text_encoder step. + negative_prompt_embeds (`Tensor`, *optional*): + negative text embeddings used to guide the image generation. Can be generated from text_encoder step. + **denoiser_input_fields (`None`, *optional*): + conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. + + Outputs: + segment_frames (`list`): + Per-segment decoded frames on CPU, each `[1, 3, T, H, W]`; the decode step concatenates, trims, and crops + them into the final video + """ + + model_name = "wan-animate-2" + block_classes = WanAnimate2CoreDenoiseBlocks.values() + block_names = WanAnimate2CoreDenoiseBlocks.keys() + + @property + def description(self): + return ( + "Core denoise step that computes the segment-invariant geometry and runs the segment-by-segment " + "denoising loop, decoding each segment inside the loop because the next segment conditions on its " + "decoded pixels." + ) + + @property + def outputs(self): + return [ + OutputParam( + "segment_frames", + type_hint=list[torch.Tensor], + description="Per-segment decoded frames on CPU, each `[1, 3, T, H, W]`; the decode step " + "concatenates, trims, and crops them into the final video", + ), + ] + + +# ==================== +# 3. Blocks +# ==================== + + +BLOCKS = InsertableDict( + [ + ("text_encoder", WanAnimate2TextEncoderStep()), + ("image_encoder", WanAnimate2ImageEncodeStep()), + ("video_encoder", WanAnimate2VideoEncodeStep()), + ("vae_encoder", WanAnimate2ImageVaeEncoderStep()), + ("denoise", WanAnimate2CoreDenoiseStep()), + ("decode", WanAnimate2DecodeStep()), + ] +) + + +# auto_docstring +class WanAnimate2Blocks(SequentialPipelineBlocks): + """ + Modular pipeline blocks for Wan-Animate-2 character animation: a reference character image and a driving video + produce a video of the character following the driving motion. + + Components: + text_encoder (`UMT5EncoderModel`) tokenizer (`AutoTokenizer`) image_processor (`WanAnimate2VideoProcessor`) + image_encoder (`CLIPVisionModel`) video_processor (`WanAnimate2VideoProcessor`) vae (`AutoencoderKLWan`) + transformer (`WanAnimate2Transformer3DModel`) scheduler (`SchedulerMixin`) guider (`ClassifierFreeGuidance`) + + Inputs: + prompt (`str`): + The prompt or prompts to guide image generation. + negative_prompt (`str`, *optional*): + The prompt or prompts not to guide the image generation. + prompt_ref (`str`, *optional*, defaults to 人物动作的参考视频): + The reference prompt for the driving video context + max_sequence_length (`int`, *optional*, defaults to 512): + Maximum sequence length for prompt encoding. + image (`Image | list`): + The reference image holding the character to animate. + height (`int`, *optional*, defaults to 800): + Together with `width`, the target *area* of the generated video; the aspect ratio comes from `image`. + Overwritten with the resolved frame height. + width (`int`, *optional*, defaults to 640): + See `height`. Overwritten with the resolved frame width. + driving_video (`list`): + The driving video that provides the motion, in any format accepted by `VideoProcessor.preprocess_video`. + driving_video_fps (`float`, *optional*): + The frame rate `driving_video` was captured at — `load_video(..., return_fps=True)` reports it. When set, + the driving frames are resampled from it to `fps`; when `None` they are used as-is. + fps (`int`, *optional*, defaults to 24): + The frame rate the model generates at + segment_frame_length (`int`, *optional*, defaults to 81): + The number of frames in each inference segment + prev_segment_conditioning_frames (`int`, *optional*, defaults to 1): + The number of conditioning frames carried over from the previous segment + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + num_inference_steps (`int`, *optional*, defaults to 40): + The number of denoising steps. + **denoiser_input_fields (`None`, *optional*): + conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. + output_type (`str`, *optional*, defaults to np): + The output type of the decoded videos + + Outputs: + videos (`list`): + The generated videos. + """ + + model_name = "wan-animate-2" + block_classes = BLOCKS.values() + block_names = BLOCKS.keys() + + @property + def description(self): + return ( + "Modular pipeline blocks for Wan-Animate-2 character animation: a reference character image and a " + "driving video produce a video of the character following the driving motion." + ) + + @property + def outputs(self): + return [OutputParam.template("videos")] diff --git a/src/diffusers/modular_pipelines/wan_animate_2/modular_blocks_wan_animate_2_distilled.py b/src/diffusers/modular_pipelines/wan_animate_2/modular_blocks_wan_animate_2_distilled.py new file mode 100644 index 000000000000..8eab815897da --- /dev/null +++ b/src/diffusers/modular_pipelines/wan_animate_2/modular_blocks_wan_animate_2_distilled.py @@ -0,0 +1,318 @@ +# 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. + +import torch + +from ...utils import logging +from ..modular_pipeline import SequentialPipelineBlocks +from ..modular_pipeline_utils import InputParam, InsertableDict, OutputParam +from .before_denoise import WanAnimate2PrepareSegmentsStep +from .decoders import WanAnimate2DecodeStep +from .denoise import WanAnimate2DistilledDenoiseStep +from .encoders import ( + WanAnimate2ImageClipEncoderStep, + WanAnimate2ImageVaeEncoderStep, + WanAnimate2ProcessImagesInputStep, + WanAnimate2ProcessVideosInputStep, + WanAnimate2TextEncoderStep, + WanAnimate2VideoClipEncoderStep, +) + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +# ==================== +# 1. Encoder groups +# ==================== + + +WanAnimate2DistilledImageEncoderBlocks = InsertableDict( + [ + ("preprocess", WanAnimate2ProcessImagesInputStep()), + ("encode", WanAnimate2ImageClipEncoderStep()), + ] +) + + +# auto_docstring +class WanAnimate2DistilledImageEncodeStep(SequentialPipelineBlocks): + """ + Image encoder step that letterboxes the reference character image to the resolved resolution and CLIP-encodes it + into `encoder_hidden_states_image`. + + Components: + image_processor (`WanAnimate2VideoProcessor`) image_encoder (`CLIPVisionModel`) + + Inputs: + image (`Image | list`): + The reference image holding the character to animate. + height (`int`, *optional*, defaults to 800): + Together with `width`, the target *area* of the generated video; the aspect ratio comes from `image`. + Overwritten with the resolved frame height. + width (`int`, *optional*, defaults to 640): + See `height`. Overwritten with the resolved frame width. + + Outputs: + image_pixels (`Tensor`): + The letterboxed reference image as a `[1, 3, H, W]` tensor in `[-1, 1]` + crop_region (`tuple`): + `(top, left, height, width)` of the reference image content inside the letterboxed frame + encoder_hidden_states_image (`Tensor`): + CLIP vision features of the reference image, conditioning every denoising forward + """ + + model_name = "wan-animate-2-distilled" + block_classes = WanAnimate2DistilledImageEncoderBlocks.values() + block_names = WanAnimate2DistilledImageEncoderBlocks.keys() + + @property + def description(self): + return ( + "Image encoder step that letterboxes the reference character image to the resolved resolution and " + "CLIP-encodes it into `encoder_hidden_states_image`." + ) + + +WanAnimate2DistilledVideoEncoderBlocks = InsertableDict( + [ + ("preprocess", WanAnimate2ProcessVideosInputStep()), + ("encode", WanAnimate2VideoClipEncoderStep()), + ] +) + + +# auto_docstring +class WanAnimate2DistilledVideoEncodeStep(SequentialPipelineBlocks): + """ + Video encoder step that preprocesses the driving video (fps resample, letterbox, zigzag padding to a whole number + of segments) and CLIP-encodes its first frame into `condition_clip_context`. + + Components: + video_processor (`WanAnimate2VideoProcessor`) image_encoder (`CLIPVisionModel`) + + Inputs: + driving_video (`list`): + The driving video that provides the motion, in any format accepted by `VideoProcessor.preprocess_video`. + driving_video_fps (`float`, *optional*): + The frame rate `driving_video` was captured at — `load_video(..., return_fps=True)` reports it. When set, + the driving frames are resampled from it to `fps`; when `None` they are used as-is. + fps (`int`, *optional*, defaults to 24): + The frame rate the model generates at + segment_frame_length (`int`, *optional*, defaults to 81): + The number of frames in each inference segment + prev_segment_conditioning_frames (`int`, *optional*, defaults to 1): + The number of conditioning frames carried over from the previous segment + height (`int`, *optional*, defaults to 800): + The height the driving frames are letterboxed to; must match the reference image's resolved height. In + the assembled pipeline the image preprocess step supplies the resolved value. + width (`int`, *optional*, defaults to 640): + See `height`. + + Outputs: + driving_video_pixels (`Tensor`): + The resampled, letterboxed, and zigzag-padded driving video, `[1, 3, T, height, width]` in `[-1, 1]` + real_frame_len (`int`): + Number of driving frames before zigzag padding; the output is trimmed to it + num_segments (`int`): + Number of inference segments + effective_segment (`int`): + Frames each segment advances by (`segment_frame_length - prev_segment_conditioning_frames`) + condition_clip_context (`Tensor`): + CLIP vision features of the driving video's first frame + """ + + model_name = "wan-animate-2-distilled" + block_classes = WanAnimate2DistilledVideoEncoderBlocks.values() + block_names = WanAnimate2DistilledVideoEncoderBlocks.keys() + + @property + def description(self): + return ( + "Video encoder step that preprocesses the driving video (fps resample, letterbox, zigzag padding to a " + "whole number of segments) and CLIP-encodes its first frame into `condition_clip_context`." + ) + + +# ==================== +# 2. Core denoise +# ==================== + + +WanAnimate2DistilledCoreDenoiseBlocks = InsertableDict( + [ + ("prepare_segments", WanAnimate2PrepareSegmentsStep()), + ("denoise", WanAnimate2DistilledDenoiseStep()), + ] +) + + +# auto_docstring +class WanAnimate2DistilledCoreDenoiseStep(SequentialPipelineBlocks): + """ + Core denoise step for the distilled Wan-Animate-2 checkpoint: computes the segment-invariant geometry and runs the + segment-by-segment denoising loop in few steps without classifier-free guidance. + + Components: + vae (`AutoencoderKLWan`) transformer (`WanAnimate2Transformer3DModel`) scheduler (`SchedulerMixin`) guider + (`ClassifierFreeGuidance`) + + Inputs: + segment_frame_length (`int`, *optional*, defaults to 81): + The number of frames in each inference segment + reference_image_latents (`Tensor`): + The reference conditioning tensor `[20, 1, latent_height, latent_width]`; provides the latent grid + driving_video_pixels (`Tensor`): + The preprocessed driving video `[1, 3, T, H, W]`, from the video preprocess step + num_segments (`int`): + Total number of segments in the driving video, from the video preprocess step + effective_segment (`int`): + Frames each segment advances: `segment_frame_length - prev_segment_conditioning_frames`, from the video + preprocess step + prev_segment_conditioning_frames (`int`, *optional*, defaults to 1): + The number of conditioning frames carried over from the previous segment + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + num_inference_steps (`int`, *optional*, defaults to 10): + The number of denoising steps. + condition_clip_context (`Tensor`): + CLIP vision features of the driving video's first frame + prompt_ref_embeds (`Tensor`): + Text embeddings of the reference prompt, guiding the reference-extraction pass + height (`int`): + The resolved frame height in pixels + width (`int`): + The resolved frame width in pixels + prompt_embeds (`Tensor`): + text embeddings used to guide the image generation. Can be generated from text_encoder step. + negative_prompt_embeds (`Tensor`, *optional*): + negative text embeddings used to guide the image generation. Can be generated from text_encoder step. + **denoiser_input_fields (`None`, *optional*): + conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. + + Outputs: + segment_frames (`list`): + Per-segment decoded frames on CPU, each `[1, 3, T, H, W]`; the decode step concatenates, trims, and crops + them into the final video + """ + + model_name = "wan-animate-2-distilled" + block_classes = WanAnimate2DistilledCoreDenoiseBlocks.values() + block_names = WanAnimate2DistilledCoreDenoiseBlocks.keys() + + @property + def description(self): + return ( + "Core denoise step for the distilled Wan-Animate-2 checkpoint: computes the segment-invariant geometry " + "and runs the segment-by-segment denoising loop in few steps without classifier-free guidance." + ) + + @property + def inputs(self): + # The distilled checkpoint samples in few steps. + return [ + InputParam.template("num_inference_steps", default=10) if param.name == "num_inference_steps" else param + for param in super().inputs + ] + + @property + def outputs(self): + return [ + OutputParam( + "segment_frames", + type_hint=list[torch.Tensor], + description="Per-segment decoded frames on CPU, each `[1, 3, T, H, W]`; the decode step " + "concatenates, trims, and crops them into the final video", + ), + ] + + +DISTILLED_BLOCKS = InsertableDict( + [ + ("text_encoder", WanAnimate2TextEncoderStep()), + ("image_encoder", WanAnimate2DistilledImageEncodeStep()), + ("video_encoder", WanAnimate2DistilledVideoEncodeStep()), + ("vae_encoder", WanAnimate2ImageVaeEncoderStep()), + ("denoise", WanAnimate2DistilledCoreDenoiseStep()), + ("decode", WanAnimate2DecodeStep()), + ] +) + + +# auto_docstring +class WanAnimate2DistilledBlocks(SequentialPipelineBlocks): + """ + Modular pipeline blocks for distilled Wan-Animate-2 character animation, sampling in few steps without + classifier-free guidance. + + Components: + text_encoder (`UMT5EncoderModel`) tokenizer (`AutoTokenizer`) image_processor (`WanAnimate2VideoProcessor`) + image_encoder (`CLIPVisionModel`) video_processor (`WanAnimate2VideoProcessor`) vae (`AutoencoderKLWan`) + transformer (`WanAnimate2Transformer3DModel`) scheduler (`SchedulerMixin`) guider (`ClassifierFreeGuidance`) + + Inputs: + prompt (`str`): + The prompt or prompts to guide image generation. + negative_prompt (`str`, *optional*): + The prompt or prompts not to guide the image generation. + prompt_ref (`str`, *optional*, defaults to 人物动作的参考视频): + The reference prompt for the driving video context + max_sequence_length (`int`, *optional*, defaults to 512): + Maximum sequence length for prompt encoding. + image (`Image | list`): + The reference image holding the character to animate. + height (`int`, *optional*, defaults to 800): + Together with `width`, the target *area* of the generated video; the aspect ratio comes from `image`. + Overwritten with the resolved frame height. + width (`int`, *optional*, defaults to 640): + See `height`. Overwritten with the resolved frame width. + driving_video (`list`): + The driving video that provides the motion, in any format accepted by `VideoProcessor.preprocess_video`. + driving_video_fps (`float`, *optional*): + The frame rate `driving_video` was captured at — `load_video(..., return_fps=True)` reports it. When set, + the driving frames are resampled from it to `fps`; when `None` they are used as-is. + fps (`int`, *optional*, defaults to 24): + The frame rate the model generates at + segment_frame_length (`int`, *optional*, defaults to 81): + The number of frames in each inference segment + prev_segment_conditioning_frames (`int`, *optional*, defaults to 1): + The number of conditioning frames carried over from the previous segment + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + num_inference_steps (`int`, *optional*, defaults to 10): + The number of denoising steps. + **denoiser_input_fields (`None`, *optional*): + conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. + output_type (`str`, *optional*, defaults to np): + The output type of the decoded videos + + Outputs: + videos (`list`): + The generated videos. + """ + + model_name = "wan-animate-2-distilled" + block_classes = DISTILLED_BLOCKS.values() + block_names = DISTILLED_BLOCKS.keys() + + @property + def description(self): + return ( + "Modular pipeline blocks for distilled Wan-Animate-2 character animation, sampling in few steps " + "without classifier-free guidance." + ) + + @property + def outputs(self): + return [OutputParam.template("videos")] diff --git a/src/diffusers/modular_pipelines/wan_animate_2/modular_pipeline.py b/src/diffusers/modular_pipelines/wan_animate_2/modular_pipeline.py new file mode 100644 index 000000000000..32ce9b222fbd --- /dev/null +++ b/src/diffusers/modular_pipelines/wan_animate_2/modular_pipeline.py @@ -0,0 +1,71 @@ +# 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. + +from ...loaders import WanLoraLoaderMixin +from ...utils import logging +from ..modular_pipeline import ModularPipeline + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +class WanAnimate2ModularPipeline(ModularPipeline, WanLoraLoaderMixin): + """ + A ModularPipeline for Wan-Animate-2 character animation. + + > [!WARNING] > This is an experimental feature and is likely to change in the future. + """ + + default_blocks_name = "WanAnimate2Blocks" + + @property + def vae_scale_factor_spatial(self): + vae_scale_factor = 8 + if hasattr(self, "vae") and self.vae is not None: + vae_scale_factor = 2 ** len(self.vae.temperal_downsample) + return vae_scale_factor + + @property + def vae_scale_factor_temporal(self): + vae_scale_factor = 4 + if hasattr(self, "vae") and self.vae is not None: + vae_scale_factor = 2 ** sum(self.vae.temperal_downsample) + return vae_scale_factor + + @property + def num_channels_latents(self): + num_channels_latents = 16 + if hasattr(self, "vae") and self.vae is not None: + num_channels_latents = self.vae.config.z_dim + return num_channels_latents + + @property + def requires_unconditional_embeds(self): + requires_unconditional_embeds = False + + if hasattr(self, "guider") and self.guider is not None: + requires_unconditional_embeds = self.guider._enabled and self.guider.num_conditions > 1 + + return requires_unconditional_embeds + + +class WanAnimate2DistilledModularPipeline(WanAnimate2ModularPipeline): + """ + A ModularPipeline for the distilled Wan-Animate-2 model, which samples in few steps without classifier-free + guidance. + + > [!WARNING] > This is an experimental feature and is likely to change in the future. + """ + + default_blocks_name = "WanAnimate2DistilledBlocks" diff --git a/src/diffusers/modular_pipelines/wan_animate_2/video_processor.py b/src/diffusers/modular_pipelines/wan_animate_2/video_processor.py new file mode 100644 index 000000000000..18e041b35db7 --- /dev/null +++ b/src/diffusers/modular_pipelines/wan_animate_2/video_processor.py @@ -0,0 +1,128 @@ +# 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. + +import numpy as np +import PIL +import torch + +from ...configuration_utils import register_to_config +from ...utils import PIL_INTERPOLATION +from ...video_processor import VideoProcessor + + +class WanAnimate2VideoProcessor(VideoProcessor): + r""" + Letterbox processor for Wan-Animate-2: `preprocess` / `preprocess_video` with `resize_mode="fill"` keep the aspect + ratio and fill the remainder with `fill_color` (black by default). The resized content is pasted at `((height - + src_h) // 2, (width - src_w) // 2)` -- the placement convention of the reference implementation -- instead of + `VaeImageProcessor`'s `(height // 2 - src_h // 2, ...)`. The two differ by one row or column whenever the frame + dimension is even and the content dimension odd, and that one-pixel placement shift is a real difference in what + the model sees. + """ + + @register_to_config + def __init__( + self, + do_resize: bool = True, + vae_scale_factor: int = 8, + vae_latent_channels: int = 16, + spatial_patch_size: tuple[int, int] = (2, 2), + resample: str = "lanczos", + reducing_gap: int = None, + do_normalize: bool = True, + do_binarize: bool = False, + do_convert_rgb: bool = False, + do_convert_grayscale: bool = False, + fill_color: str | float | tuple[float, ...] | None = 0, + ): + # Deliberately does not chain into the parent `__init__`: it is itself + # `register_to_config`-decorated, so calling it re-registers every shared field with its + # defaults and discards what the caller asked for -- `resample` included. The decorator on + # this method already records the full config; the parent's body only validates. + if do_convert_rgb and do_convert_grayscale: + raise ValueError( + "`do_convert_rgb` and `do_convert_grayscale` can not both be set to `True`," + " if you intended to convert the image into RGB format, please set `do_convert_grayscale = False`.", + " if you intended to convert the image into grayscale format, please set `do_convert_rgb = False`", + ) + + def _resize_and_fill( + self, + image: PIL.Image.Image, + width: int, + height: int, + ) -> PIL.Image.Image: + ratio = width / height + src_ratio = image.width / image.height + + src_w = width if ratio < src_ratio else image.width * height // image.height + src_h = height if ratio >= src_ratio else image.height * width // image.width + + resized = image.resize((src_w, src_h), resample=PIL_INTERPOLATION[self.config.resample]) + res = PIL.Image.new("RGB", (width, height), color=self.config.fill_color or 0) + res.paste(resized, box=((width - src_w) // 2, (height - src_h) // 2)) + return res + + # Copied from diffusers.pipelines.wan.image_processor.WanAnimateImageProcessor.get_default_height_width + def get_default_height_width( + self, + image: PIL.Image.Image | np.ndarray | torch.Tensor, + height: int | None = None, + width: int | None = None, + ) -> tuple[int, int]: + r""" + Returns the height and width of the image, downscaled to the next integer multiple of `vae_scale_factor`. + + Args: + image (`PIL.Image.Image | np.ndarray | torch.Tensor`): + The image input, which can be a PIL image, NumPy array, or PyTorch tensor. If it is a NumPy array, it + should have shape `[batch, height, width]` or `[batch, height, width, channels]`. If it is a PyTorch + tensor, it should have shape `[batch, channels, height, width]`. + height (`int | None`, *optional*, defaults to `None`): + The height of the preprocessed image. If `None`, the height of the `image` input will be used. + width (`int | None`, *optional*, defaults to `None`): + The width of the preprocessed image. If `None`, the width of the `image` input will be used. + + Returns: + `tuple[int, int]`: + A tuple containing the height and width, both resized to the nearest integer multiple of + `vae_scale_factor * spatial_patch_size`. + """ + + if height is None: + if isinstance(image, PIL.Image.Image): + height = image.height + elif isinstance(image, torch.Tensor): + height = image.shape[2] + else: + height = image.shape[1] + + if width is None: + if isinstance(image, PIL.Image.Image): + width = image.width + elif isinstance(image, torch.Tensor): + width = image.shape[3] + else: + width = image.shape[2] + + max_area = width * height + aspect_ratio = height / width + mod_value_h = self.config.vae_scale_factor * self.config.spatial_patch_size[0] + mod_value_w = self.config.vae_scale_factor * self.config.spatial_patch_size[1] + + # Try to preserve the aspect ratio + height = round(np.sqrt(max_area * aspect_ratio)) // mod_value_h * mod_value_h + width = round(np.sqrt(max_area / aspect_ratio)) // mod_value_w * mod_value_w + + return height, width diff --git a/src/diffusers/utils/dummy_pt_objects.py b/src/diffusers/utils/dummy_pt_objects.py index 0e17254611c8..b8c370bfd1e6 100644 --- a/src/diffusers/utils/dummy_pt_objects.py +++ b/src/diffusers/utils/dummy_pt_objects.py @@ -2310,6 +2310,21 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch"]) +class WanAnimate2Transformer3DModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + class WanAnimateTransformer3DModel(metaclass=DummyObject): _backends = ["torch"] diff --git a/src/diffusers/utils/dummy_torch_and_transformers_objects.py b/src/diffusers/utils/dummy_torch_and_transformers_objects.py index 907bff826d11..11e69fcdccf1 100644 --- a/src/diffusers/utils/dummy_torch_and_transformers_objects.py +++ b/src/diffusers/utils/dummy_torch_and_transformers_objects.py @@ -842,6 +842,66 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch", "transformers"]) +class WanAnimate2Blocks(metaclass=DummyObject): + _backends = ["torch", "transformers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch", "transformers"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + +class WanAnimate2DistilledBlocks(metaclass=DummyObject): + _backends = ["torch", "transformers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch", "transformers"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + +class WanAnimate2DistilledModularPipeline(metaclass=DummyObject): + _backends = ["torch", "transformers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch", "transformers"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + +class WanAnimate2ModularPipeline(metaclass=DummyObject): + _backends = ["torch", "transformers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch", "transformers"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + class WanBlocks(metaclass=DummyObject): _backends = ["torch", "transformers"] diff --git a/src/diffusers/utils/loading_utils.py b/src/diffusers/utils/loading_utils.py index c4fee0cfdd83..ba528c7f857c 100644 --- a/src/diffusers/utils/loading_utils.py +++ b/src/diffusers/utils/loading_utils.py @@ -57,7 +57,8 @@ def load_image( def load_video( video: str, convert_method: Callable[[list[PIL.Image.Image]], list[PIL.Image.Image]] | None = None, -) -> list[PIL.Image.Image]: + return_fps: bool = False, +) -> list[PIL.Image.Image] | tuple[list[PIL.Image.Image], float]: """ Loads `video` to a list of PIL Image. @@ -67,10 +68,13 @@ def load_video( convert_method (Callable[[list[PIL.Image.Image]], list[PIL.Image.Image]], *optional*): A conversion method to apply to the video after loading it. When set to `None` the images will be converted to "RGB". + return_fps (`bool`, *optional*, defaults to `False`): + Whether to also return the frame rate the video was encoded at. Needed by pipelines that resample the input + to the frame rate their model works at, since a list of frames does not carry that information. Returns: - `list[PIL.Image.Image]`: - The video as a list of PIL images. + `list[PIL.Image.Image]` or `tuple[list[PIL.Image.Image], float]`: + The video as a list of PIL images, and its frame rate if `return_fps` is set. """ is_url = video.startswith("http://") or video.startswith("https://") is_file = os.path.isfile(video) @@ -102,8 +106,12 @@ def load_video( video = video_path pil_images = [] + fps = None if video.endswith(".gif"): gif = PIL.Image.open(video) + # Milliseconds this frame is displayed for; GIFs are not obliged to record it. + frame_duration = gif.info.get("duration") + fps = 1000 / frame_duration if frame_duration else None try: while True: pil_images.append(gif.copy()) @@ -125,6 +133,7 @@ def load_video( ) with imageio.get_reader(video) as reader: + fps = reader.get_meta_data().get("fps") # Read all frames for frame in reader: pil_images.append(PIL.Image.fromarray(frame)) @@ -135,6 +144,11 @@ def load_video( if convert_method is not None: pil_images = convert_method(pil_images) + if return_fps: + if fps is None: + raise ValueError(f"Could not read a frame rate from {video}, so `return_fps=True` cannot be honoured.") + return pil_images, fps + return pil_images diff --git a/tests/modular_pipelines/test_modular_pipelines_common.py b/tests/modular_pipelines/test_modular_pipelines_common.py index fadaef9d3e58..759da1ac13b0 100644 --- a/tests/modular_pipelines/test_modular_pipelines_common.py +++ b/tests/modular_pipelines/test_modular_pipelines_common.py @@ -532,7 +532,8 @@ def test_workflow_defaults(self): blocks = self.pipeline_blocks_class() for workflow_name, expected_defaults in self.expected_workflow_defaults.items(): - workflow_blocks = blocks.get_workflow(workflow_name) + # a pipeline without workflows is tested as the single unnamed workflow `None` over the full blockset + workflow_blocks = blocks if workflow_name is None else blocks.get_workflow(workflow_name) # components: every one of the workflow is named with its class, so one appearing, disappearing or # changing type fails loudly @@ -577,6 +578,17 @@ def test_workflow_defaults(self): f"{param.default!r} (required={param.required}), expected {expected_default!r}" ) + # component configs, optionally: pinned config values for the `from_config` components worth + # watching — e.g. a preset's guider scale — checked against the spec that creates them. Pretrained + # components take their config from the repo, so there is nothing block-level to pin. + for component_name, expected_config in expected_defaults.get("component_configs", {}).items(): + actual_config = dict(component_specs[component_name].config or {}) + for config_name, expected_value in expected_config.items(): + assert actual_config.get(config_name) == expected_value, ( + f"Workflow '{workflow_name}': component '{component_name}' config " + f"'{config_name}' is {actual_config.get(config_name)!r}, expected {expected_value!r}" + ) + def test_from_pretrained_workflow(self): blocks = self.pipeline_blocks_class() if blocks._workflow_map is None: diff --git a/tests/modular_pipelines/wan_animate_2/__init__.py b/tests/modular_pipelines/wan_animate_2/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tests/modular_pipelines/wan_animate_2/test_modular_pipeline_wan_animate_2.py b/tests/modular_pipelines/wan_animate_2/test_modular_pipeline_wan_animate_2.py new file mode 100644 index 000000000000..44e636de0e94 --- /dev/null +++ b/tests/modular_pipelines/wan_animate_2/test_modular_pipeline_wan_animate_2.py @@ -0,0 +1,133 @@ +# coding=utf-8 +# Copyright 2026 HuggingFace Inc. +# +# 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. + +import numpy as np +import PIL.Image +import pytest + +from diffusers.modular_pipelines import ( + WanAnimate2Blocks, + WanAnimate2DistilledBlocks, + WanAnimate2DistilledModularPipeline, + WanAnimate2ModularPipeline, +) + +from ..test_modular_pipelines_common import ModularGuiderTesterMixin, ModularPipelineTesterMixin + + +# Every component with its class, every input — optional ones with their defaults — and the config values worth +# watching: the guider scale is what tells the two presets apart. +_COMPONENTS = { + "text_encoder": "UMT5EncoderModel", + "tokenizer": "AutoTokenizer", + "image_processor": "WanAnimate2VideoProcessor", + "image_encoder": "CLIPVisionModel", + "video_processor": "WanAnimate2VideoProcessor", + "vae": "AutoencoderKLWan", + "transformer": "WanAnimate2Transformer3DModel", + "scheduler": "SchedulerMixin", + "guider": "ClassifierFreeGuidance", +} +_INPUTS = { + "negative_prompt": None, + "prompt_ref": "人物动作的参考视频", + "max_sequence_length": 512, + "height": 800, + "width": 640, + "driving_video_fps": None, + "fps": 24, + "segment_frame_length": 81, + "prev_segment_conditioning_frames": 1, + "generator": None, + "output_type": "np", +} +WAN_ANIMATE_2_DEFAULTS = { + None: { + "components": _COMPONENTS, + "configs": {}, + "required_inputs": ["prompt", "image", "driving_video"], + "inputs": {**_INPUTS, "num_inference_steps": 40}, + "component_configs": {"guider": {"guidance_scale": 3.0}}, + } +} +WAN_ANIMATE_2_DISTILLED_DEFAULTS = { + None: { + "components": _COMPONENTS, + "configs": {}, + "required_inputs": ["prompt", "image", "driving_video"], + "inputs": {**_INPUTS, "num_inference_steps": 10}, + "component_configs": {"guider": {"guidance_scale": 1.0}}, + } +} + + +class TestWanAnimate2ModularPipelineFast(ModularPipelineTesterMixin, ModularGuiderTesterMixin): + pipeline_class = WanAnimate2ModularPipeline + pipeline_blocks_class = WanAnimate2Blocks + pretrained_model_name_or_path = "hf-internal-testing/tiny-wan-animate-2-modular" + + params = frozenset(["prompt", "image", "driving_video"]) + batch_params = frozenset() + optional_params = frozenset(["num_inference_steps", "height", "width", "output_type"]) + output_name = "videos" + expected_workflow_defaults = WAN_ANIMATE_2_DEFAULTS + + def get_dummy_inputs(self, seed=0): + rng = np.random.RandomState(seed) + image = PIL.Image.fromarray(rng.randint(0, 255, (96, 64, 3), dtype=np.uint8)) + driving_video = [PIL.Image.fromarray(rng.randint(0, 255, (40, 32, 3), dtype=np.uint8)) for _ in range(17)] + inputs = { + "prompt": "a tiny cat", + "image": image, + "driving_video": driving_video, + "generator": self.get_generator(seed), + "num_inference_steps": 2, + "height": 64, + "width": 64, + "segment_frame_length": 9, + "prev_segment_conditioning_frames": 1, + "max_sequence_length": 16, + "output_type": "pt", + } + return inputs + + @pytest.mark.skip(reason="Wan-Animate-2 is unbatched: one character image and driving video per call") + def test_inference_batch_consistent(self): + pass + + @pytest.mark.skip(reason="Wan-Animate-2 is unbatched: one character image and driving video per call") + def test_inference_batch_single_identical(self): + pass + + @pytest.mark.skip(reason="Wan-Animate-2 is unbatched: no num_videos_per_prompt") + def test_num_images_per_prompt(self): + pass + + def test_guider_cfg(self): + # The tiny random transformer responds only weakly to the text embeddings (cond vs uncond + # predictions differ by ~1e-4), so the CFG effect on the decoded pixels is real but small. + # Same-seed runs are bit-deterministic, so any nonzero difference here is guidance. + super().test_guider_cfg(expected_max_diff=1e-6) + + +class TestWanAnimate2DistilledModularPipelineFast(TestWanAnimate2ModularPipelineFast): + pipeline_class = WanAnimate2DistilledModularPipeline + pipeline_blocks_class = WanAnimate2DistilledBlocks + pretrained_model_name_or_path = "hf-internal-testing/tiny-wan-animate-2-distilled-modular" + expected_workflow_defaults = WAN_ANIMATE_2_DISTILLED_DEFAULTS + + @pytest.mark.skip(reason="The distilled preset pins its guider to guidance_scale=1.0") + def test_guider_cfg(self): + pass