From 636560149dec91488557682302e620d452a91a5a Mon Sep 17 00:00:00 2001 From: kelseyee <971704395@qq.com> Date: Thu, 30 Jul 2026 15:59:32 +0800 Subject: [PATCH 01/19] support wan-animate-2 --- src/diffusers/__init__.py | 4 + src/diffusers/loaders/single_file_model.py | 5 + src/diffusers/loaders/single_file_utils.py | 21 + src/diffusers/models/__init__.py | 2 + src/diffusers/models/transformers/__init__.py | 1 + .../transformers/transformer_wan_animate_2.py | 1120 +++++++++++++++++ .../modular_pipelines/wan/__init__.py | 4 + .../wan/modular_blocks_wan_animate_2.py | 462 +++++++ .../modular_pipelines/wan/modular_pipeline.py | 10 + src/diffusers/pipelines/__init__.py | 2 + src/diffusers/pipelines/pipeline_utils.py | 24 +- src/diffusers/pipelines/wan/__init__.py | 2 + .../pipelines/wan/pipeline_wan_animate_2.py | 731 +++++++++++ src/diffusers/utils/dummy_pt_objects.py | 15 + .../dummy_torch_and_transformers_objects.py | 15 + 15 files changed, 2413 insertions(+), 5 deletions(-) create mode 100644 src/diffusers/models/transformers/transformer_wan_animate_2.py create mode 100644 src/diffusers/modular_pipelines/wan/modular_blocks_wan_animate_2.py create mode 100644 src/diffusers/pipelines/wan/pipeline_wan_animate_2.py diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py index 7a8d727aefea..b6e22e7891fd 100644 --- a/src/diffusers/__init__.py +++ b/src/diffusers/__init__.py @@ -349,6 +349,7 @@ "UVit2DModel", "VQModel", "WanAnimateTransformer3DModel", + "WanAnimate2Transformer3DModel", "WanTransformer3DModel", "WanVACETransformer3DModel", "ZImageControlNetModel", @@ -838,6 +839,7 @@ "VisualClozePipeline", "VQDiffusionPipeline", "WanAnimatePipeline", + "WanAnimate2Pipeline", "WanImageToVideoPipeline", "WanPipeline", "WanVACEPipeline", @@ -1188,6 +1190,7 @@ UNetSpatioTemporalConditionModel, UVit2DModel, VQModel, + WanAnimate2Transformer3DModel, WanAnimateTransformer3DModel, WanTransformer3DModel, WanVACETransformer3DModel, @@ -1651,6 +1654,7 @@ VisualClozeGenerationPipeline, VisualClozePipeline, VQDiffusionPipeline, + WanAnimate2Pipeline, WanAnimatePipeline, WanImageToVideoPipeline, WanPipeline, 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..c22ddb9a3a18 100644 --- a/src/diffusers/loaders/single_file_utils.py +++ b/src/diffusers/loaders/single_file_utils.py @@ -3289,6 +3289,27 @@ 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. + """ + converted_state_dict = {} + + # Strip model.diffusion_model prefix if present + keys = list(checkpoint.keys()) + for k in keys: + if "model.diffusion_model." in k: + checkpoint[k.replace("model.diffusion_model.", "")] = checkpoint.pop(k) + + # The official checkpoint already uses the same key format as the diffusers model + # (blocks.N.block.*), so no remapping is needed. + for key in list(checkpoint.keys()): + converted_state_dict[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 167ee7a534de..b82d2f4d6d39 100755 --- a/src/diffusers/models/__init__.py +++ b/src/diffusers/models/__init__.py @@ -141,6 +141,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"] @@ -278,6 +279,7 @@ Transformer2DModel, TransformerTemporalModel, WanAnimateTransformer3DModel, + WanAnimate2Transformer3DModel, WanTransformer3DModel, WanVACETransformer3DModel, ZImageTransformer2DModel, diff --git a/src/diffusers/models/transformers/__init__.py b/src/diffusers/models/transformers/__init__.py index 21f5cb853643..ad433e6edc09 100755 --- a/src/diffusers/models/transformers/__init__.py +++ b/src/diffusers/models/transformers/__init__.py @@ -63,5 +63,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..e9edd44b1f9c --- /dev/null +++ b/src/diffusers/models/transformers/transformer_wan_animate_2.py @@ -0,0 +1,1120 @@ +# 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 +from functools import lru_cache, partial + +import numpy as np +import torch +import torch.nn as nn +from torch.nn.attention.flex_attention import create_block_mask + +from ...configuration_utils import ConfigMixin, register_to_config +from ...loaders import FromOriginalModelMixin, PeftAdapterMixin +from ..modeling_utils import ModelMixin + + +try: + from flash_attn_interface import flash_attn_varlen_func + + FLASH_VER = 3 +except ModuleNotFoundError: + try: + from flash_attn import flash_attn_varlen_func + + FLASH_VER = 2 + except ModuleNotFoundError: + flash_attn_varlen_func = None + FLASH_VER = None + +from torch.nn.attention.flex_attention import flex_attention as _flex_attention_raw + +# Lazy compile: compile on first call instead of at import time +_flex_compiled = None + + +def _get_compiled_flex_attention(): + global _flex_compiled + if _flex_compiled is None: + _flex_compiled = torch.compile(_flex_attention_raw, dynamic=False, mode="max-autotune", fullgraph=True) + return _flex_compiled + + +def flash_attention( + q, + k, + v, + q_lens=None, + k_lens=None, + dropout_p=0.0, + softmax_scale=None, + q_scale=None, + causal=False, + window_size=(-1, -1), + deterministic=False, + dtype=torch.bfloat16, +): + """ + q: [B, Lq, Nq, C1]. + k: [B, Lk, Nk, C1]. + v: [B, Lk, Nk, C2]. Nq must be divisible by Nk. + q_lens: [B]. + k_lens: [B]. + dropout_p: float. Dropout probability. + softmax_scale: float. The scaling of QK^T before applying softmax. + causal: bool. Whether to apply causal attention mask. + window_size: (left right). If not (-1, -1), apply sliding window local attention. + deterministic: bool. If True, slightly slower and uses more memory. + dtype: torch.dtype. Apply when dtype of q/k/v is not float16/bfloat16. + """ + half_dtypes = (torch.float16, torch.bfloat16) + assert dtype in half_dtypes + assert q.device.type == "cuda" and q.size(-1) <= 256 + + # params + b, lq, lk, out_dtype = q.size(0), q.size(1), k.size(1), q.dtype + + def half(x): + return x if x.dtype in half_dtypes else x.to(dtype) + + # preprocess query + if q_lens is None: + q = half(q.flatten(0, 1)) + q_lens = torch.tensor([lq] * b, dtype=torch.int32).to(device=q.device, non_blocking=True) + else: + q = half(torch.cat([u[:v] for u, v in zip(q, q_lens)])) + + # preprocess key, value + if k_lens is None: + k = half(k.flatten(0, 1)) + v = half(v.flatten(0, 1)) + k_lens = torch.tensor([lk] * b, dtype=torch.int32).to(device=k.device, non_blocking=True) + else: + k = half(torch.cat([u[:v] for u, v in zip(k, k_lens)])) + v = half(torch.cat([u[:v] for u, v in zip(v, k_lens)])) + + q = q.to(v.dtype) + k = k.to(v.dtype) + + if q_scale is not None: + q = q * q_scale + # apply attention + if FLASH_VER == 3: + # Note: dropout_p, window_size are not supported in FA3 now. + x = flash_attn_varlen_func( + q=q, + k=k, + v=v, + cu_seqlens_q=torch.cat([q_lens.new_zeros([1]), q_lens]) + .cumsum(0, dtype=torch.int32) + .to(q.device, non_blocking=True), + cu_seqlens_k=torch.cat([k_lens.new_zeros([1]), k_lens]) + .cumsum(0, dtype=torch.int32) + .to(q.device, non_blocking=True), + max_seqlen_q=lq, + max_seqlen_k=lk, + softmax_scale=softmax_scale, + causal=causal, + deterministic=deterministic, + )[0].unflatten(0, (b, lq)) + else: + assert FLASH_VER == 2 + x = flash_attn_varlen_func( + q=q, + k=k, + v=v, + cu_seqlens_q=torch.cat([q_lens.new_zeros([1]), q_lens]) + .cumsum(0, dtype=torch.int32) + .to(q.device, non_blocking=True), + cu_seqlens_k=torch.cat([k_lens.new_zeros([1]), k_lens]) + .cumsum(0, dtype=torch.int32) + .to(q.device, non_blocking=True), + max_seqlen_q=lq, + max_seqlen_k=lk, + dropout_p=dropout_p, + softmax_scale=softmax_scale, + causal=causal, + window_size=window_size, + deterministic=deterministic, + ).unflatten(0, (b, lq)) + + # output + return x.type(out_dtype) + + +def flex_attention( + q, + k, + v, + q_lens=None, + k_lens=None, + block_mask=None, + kernel_options=None, + dtype=torch.bfloat16, + score_mod=None, +): + """ + q: [B, Lq, Nq, C1]. + k: [B, Lk, Nk, C1]. + v: [B, Lk, Nk, C2]. Nq must be divisible by Nk. + q_lens: [B]. + k_lens: [B]. + dtype: torch.dtype. Apply when dtype of q/k/v is not float16/bfloat16. + """ + half_dtypes = (torch.float16, torch.bfloat16) + assert dtype in half_dtypes + assert q.device.type == "cuda" + lq, lk, out_dtype = q.size(1), k.size(1), q.dtype + + def half(x): + return x if x.dtype in half_dtypes else x.to(dtype) + + assert lq % 128 == 0, "q_len must be divisible by 128." + assert lk % 128 == 0, "k_len must be divisible by 128." + + # preprocess query + if q_lens is None: + q = half(q) + else: + q = half(q) + assert q_lens.max() == q_lens.min(), "varlen of query is not supported" + + # preprocess key, value + if k_lens is None: + k, v = half(k), half(v) + else: + k, v = half(k), half(v) + assert k_lens.max() == k_lens.min(), "varlen of key is not supported" + + q = q.to(v.dtype) + k = k.to(v.dtype) + + x = _get_compiled_flex_attention()( + query=q.transpose(2, 1), + key=k.transpose(2, 1), + value=v.transpose(2, 1), + block_mask=block_mask, + kernel_options=kernel_options, + score_mod=score_mod, + ).transpose(2, 1) + + return x.type(out_dtype) + + +def _score_mod_impl(score, b_idx, h_idx, q_idx, kv_idx, hw: int, log_scale: float): + condition = (kv_idx >= hw) & (kv_idx < 2 * hw) + return torch.where(condition, score + log_scale, score) + + +@lru_cache(maxsize=32) +def _get_score_mod(hw: int, log_scale: float = -1.0): + return partial(_score_mod_impl, hw=hw, log_scale=log_scale) + + +def sinusoidal_embedding_1d(dim, position): + # preprocess + assert dim % 2 == 0 + half = dim // 2 + position = position.type(torch.float64) + + # calculation + sinusoid = torch.outer(position, torch.pow(10000, -torch.arange(half).to(position).div(half))) + x = torch.cat([torch.cos(sinusoid), torch.sin(sinusoid)], dim=1) + return x + + +@torch.amp.autocast(device_type="cuda", enabled=False) +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 + + +@torch.amp.autocast(device_type="cuda", enabled=False) +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 + + +class RMSNorm(nn.Module): + def __init__(self, dim, eps=1e-5): + super().__init__() + self.dim = dim + self.eps = eps + self.weight = nn.Parameter(torch.ones(dim)) + + def forward(self, x): + return self._norm(x.float()).type_as(x) * self.weight + + def _norm(self, x): + return x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps) + + +class LayerNorm(nn.LayerNorm): + """ + LayerNorm without learnable affine parameters. + """ + + def __init__(self, dim, eps=1e-6, elementwise_affine=False): + super().__init__(dim, elementwise_affine=elementwise_affine, eps=eps) + + def forward(self, x): + return super().forward(x.float()).type_as(x) + + +class SelfAttention(nn.Module): + def __init__( + self, + dim, + num_heads, + window_size=(-1, -1), + qk_norm=True, + eps=1e-6, + ): + assert dim % num_heads == 0 + super().__init__() + self.dim = dim + self.num_heads = num_heads + self.head_dim = dim // num_heads + self.window_size = window_size + self.qk_norm = qk_norm + self.eps = eps + + # layers + self.q = nn.Linear(dim, dim) + self.k = nn.Linear(dim, dim) + self.v = nn.Linear(dim, dim) + self.o = nn.Linear(dim, dim) + self.norm_q = RMSNorm(dim, eps=eps) if qk_norm else nn.Identity() + self.norm_k = RMSNorm(dim, eps=eps) if qk_norm else nn.Identity() + + def forward(self, *args, method, **kwargs): + return getattr(self, method)(*args, **kwargs) + + def pre_attention(self, x): + b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim + + # query, key, value function + def qkv_fn(x): + q = self.norm_q(self.q(x)).view(b, s, n, d) + k = self.norm_k(self.k(x)).view(b, s, n, d) + v = self.v(x).view(b, s, n, d) + return q, k, v + + q, k, v = qkv_fn(x) + + return q, k, v + + def post_attention(self, x): + # output + x = x.flatten(2) + x = self.o(x) + return x + + +class CrossAttention(SelfAttention): + def __init__(self, dim, num_heads, window_size=(-1, -1), qk_norm=True, eps=1e-6, use_img_emb=True): + super().__init__(dim, num_heads, window_size, qk_norm, eps) + self.use_img_emb = use_img_emb + if use_img_emb: + self.k_img = nn.Linear(dim, dim) + self.v_img = nn.Linear(dim, dim) + self.norm_k_img = RMSNorm(dim, eps=eps) if qk_norm else nn.Identity() + + def forward(self, x, context, context_lens, counter=0): + """ + x: [B, L1, C]. + context: [B, L2, C]. + context_lens: [B]. + """ + if self.use_img_emb: + context_img = context[:, :257] + context = context[:, 257:] + else: + context = context + + b, n, d = x.size(0), self.num_heads, self.head_dim + + # compute query, key, value + q = self.norm_q(self.q(x)).view(b, -1, n, d) + k = self.norm_k(self.k(context)).view(b, -1, n, d) + v = self.v(context).view(b, -1, n, d) + + if self.use_img_emb: + k_img = self.norm_k_img(self.k_img(context_img)).view(b, -1, n, d) + v_img = self.v_img(context_img).view(b, -1, n, d) + img_x = flash_attention(q, k_img, v_img, k_lens=None) + # compute attention + x = flash_attention(q, k, v, k_lens=context_lens) + + # output + x = x.flatten(2) + if self.use_img_emb: + img_x = img_x.flatten(2) + x = x + img_x + x = self.o(x) + return x + + +class AttentionBlock(nn.Module): + def __init__( + self, + dim, + ffn_dim, + num_heads, + window_size=(-1, -1), + qk_norm=True, + cross_attn_norm=False, + eps=1e-6, + use_img_emb=True, + ): + super().__init__() + self.dim = dim + self.ffn_dim = ffn_dim + self.num_heads = num_heads + self.window_size = window_size + self.qk_norm = qk_norm + self.cross_attn_norm = cross_attn_norm + self.eps = eps + + # layers + self.norm1 = LayerNorm(dim, eps) + + self.self_attn = SelfAttention(dim, num_heads, window_size, qk_norm, eps) + + self.norm3 = LayerNorm(dim, eps, elementwise_affine=True) if cross_attn_norm else nn.Identity() + + self.cross_attn = CrossAttention(dim, num_heads, (-1, -1), qk_norm, eps, use_img_emb=use_img_emb) + + self.norm2 = LayerNorm(dim, eps) + self.ffn = nn.Sequential( + nn.Linear(dim, ffn_dim), + nn.GELU(approximate="tanh"), + nn.Linear(ffn_dim, dim), + ) + # modulation + self.modulation = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5) + + def forward(self, *args, method, **kwargs): + return getattr(self, method)(*args, **kwargs) + + def pre_self_attention(self, x, e): + assert e.dtype == torch.float32 + with torch.amp.autocast(device_type="cuda", dtype=torch.float32): + e = (self.modulation + e).chunk(6, dim=1) + assert e[0].dtype == torch.float32 + + q, k, v = self.self_attn(self.norm1(x).float() * (1 + e[1]) + e[0], method="pre_attention") + return q, k, v, e + + def post_self_attention(self, x): + x = self.self_attn(x, method="post_attention") + return x + + def cross_attention(self, x, context, context_lens, e): + x = x + self.cross_attn(self.norm3(x), context, context_lens) + y = self.ffn(self.norm2(x).float() * (1 + e[4]) + e[3]) + with torch.amp.autocast(device_type="cuda", dtype=torch.float32): + x = x + y * e[5] + return x + + +class IncontextAttentionBlock(nn.Module): + def __init__( + self, + dim, + ffn_dim, + num_heads, + window_size=(-1, -1), + qk_norm=True, + cross_attn_norm=False, + eps=1e-6, + refer_stride=1, + use_img_emb=True, + sparse_type=0, + log_scale=0.0, + ): + super().__init__() + self.dim = dim + self.ffn_dim = ffn_dim + self.num_heads = num_heads + self.window_size = window_size + self.qk_norm = qk_norm + self.cross_attn_norm = cross_attn_norm + self.eps = eps + self.refer_stride = refer_stride + self.sparse_type = sparse_type + self.log_scale = log_scale + + self.block = AttentionBlock( + dim, ffn_dim, num_heads, window_size, qk_norm, cross_attn_norm, eps, use_img_emb=use_img_emb + ) + + def forward(self, *args, method, **kwargs): + return getattr(self, method)(*args, **kwargs) + + def forward_ref(self, x_ref, index, k_cache, v_cache, context_ref, freqs_ref, grid_sizes_ref, e_ref, context_lens): + q_ref, k_ref, v_ref, e_ref = self.block(x_ref, e_ref, method="pre_self_attention") + + k_cache[index] = k_ref + v_cache[index] = v_ref + q_ref_add_rope = rope_apply(q_ref, grid_sizes_ref, freqs_ref, self.refer_stride) + k_ref_add_rope = rope_apply(k_ref, grid_sizes_ref, freqs_ref, self.refer_stride) + + ref_f, ref_h, ref_w = grid_sizes_ref[0].tolist() + ref_vail_len = ref_f * ref_h * ref_w + + xout_ref = flash_attention( + q=q_ref_add_rope, + k=k_ref_add_rope, + v=v_ref, + k_lens=torch.tensor([ref_vail_len], dtype=torch.long), + window_size=self.window_size, + ) + + y_ref = self.block(xout_ref, method="post_self_attention") + + with torch.amp.autocast(device_type="cuda", dtype=torch.float32): + x_ref = x_ref + y_ref * e_ref[2] + + x_ref = self.block(x_ref, context_ref, context_lens, e_ref, method="cross_attention") + + return x_ref + + def forward_gen( + self, + x, + index, + k_cache, + v_cache, + block_mask, + context, + freqs, + freqs_ref, + grid_sizes, + grid_sizes_ref, + origin_len, + origin_area, + e, + context_lens, + ): + origin_latent_f = origin_len // 4 + 1 + origin_latent_hw = origin_area[0] * origin_area[1] // 256 + origin_max_len = (origin_latent_f + 1) * origin_latent_hw + origin_ref_max_len = origin_latent_f * origin_latent_hw + + f, h, w = grid_sizes[0].tolist() + vail_len = f * h * w + hw = h * w + + ref_f, ref_h, ref_w = grid_sizes_ref[0].tolist() + ref_vail_len = ref_f * ref_h * ref_w + ref_hw = ref_h * ref_w + + q, k, v, e = self.block(x, e, method="pre_self_attention") + + q = rope_apply(q, grid_sizes, freqs) + k = rope_apply(k, grid_sizes, freqs) + k_ref, v_ref = k_cache[index], v_cache[index] + k_ref = rope_apply(k_ref, grid_sizes_ref, freqs_ref, self.refer_stride) + + B, _, N, C = q.shape + device, dtype = q.device, q.dtype + + target_q_len = math.ceil(origin_max_len / 128) * 128 + target_ref_len = math.ceil(origin_ref_max_len / 128) * 128 + target_kv_len = target_q_len + target_ref_len + + q_padding = q[:, vail_len:].clone() + + q_incontext = torch.zeros(B, target_q_len, N, C, device=device, dtype=dtype) + k_incontext = torch.zeros(B, target_kv_len, N, C, device=device, dtype=dtype) + v_incontext = torch.zeros(B, target_kv_len, N, C, device=device, dtype=dtype) + + q_src = q[:, :vail_len].view(B, f, hw, N, C) + k_src = k[:, :vail_len].view(B, f, hw, N, C) + v_src = v[:, :vail_len].view(B, f, hw, N, C) + + q_incontext[:, : f * origin_latent_hw].view(B, f, origin_latent_hw, N, C)[:, :, :hw] = q_src + k_incontext[:, : f * origin_latent_hw].view(B, f, origin_latent_hw, N, C)[:, :, :hw] = k_src + v_incontext[:, : f * origin_latent_hw].view(B, f, origin_latent_hw, N, C)[:, :, :hw] = v_src + + k_ref_src = k_ref[:, :ref_vail_len].view(B, ref_f, ref_hw, N, C) + v_ref_src = v_ref[:, :ref_vail_len].view(B, ref_f, ref_hw, N, C) + + k_incontext[:, target_q_len : target_q_len + ref_f * origin_latent_hw].view(B, ref_f, origin_latent_hw, N, C)[ + :, :, :ref_hw + ] = k_ref_src + v_incontext[:, target_q_len : target_q_len + ref_f * origin_latent_hw].view(B, ref_f, origin_latent_hw, N, C)[ + :, :, :ref_hw + ] = v_ref_src + + score_mod = _get_score_mod(hw=int(origin_latent_hw), log_scale=self.log_scale) + + xout_full = flex_attention( + q=q_incontext, + k=k_incontext, + v=v_incontext, + block_mask=block_mask, + kernel_options=None, + score_mod=score_mod, + ) + + xout_valid = xout_full[:, : f * origin_latent_hw] + xout_valid = xout_valid.view(B, f, origin_latent_hw, N, C) + xout_vail = xout_valid[:, :, :hw] + xout_vail = xout_vail.reshape(B, f * hw, N, C) + xout = torch.cat([xout_vail, q_padding], dim=1) + + y = self.block(xout, method="post_self_attention") + + with torch.amp.autocast(device_type="cuda", dtype=torch.float32): + x = x + y * e[2] + + x = self.block(x, context, context_lens, e, method="cross_attention") + return x + + +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 = LayerNorm(dim, eps) + 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): + assert e.dtype == torch.float32 + with torch.amp.autocast(device_type="cuda", dtype=torch.float32): + e = (self.modulation + e.unsqueeze(1)).chunk(2, dim=1) + x = self.head(self.norm(x) * (1 + e[1]) + e[0]) + 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): + 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 KV cache: a reference video is first encoded + (``forward_ref``) to cache K/V tensors, then the generation forward (``forward_gen``) uses the cached + K/V with a block mask (``flex_attention``) and score modification (``log_scale``) for frame-level + sparse in-context attention. + + 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. + window_size (`tuple[int]`, defaults to `(-1, -1)`): + Window size for local attention (-1 indicates global attention). + qk_norm (`bool`, defaults to `True`): + Enable query/key normalization. + 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. + sparse_type (`int`, defaults to `0`): + Sparse attention type. + log_scale (`float`, defaults to `0.0`): + Log scale for score modification in in-context attention. + """ + + _supports_gradient_checkpointing = True + _skip_layerwise_casting_patterns = ["patch_embedding", "img_emb", "norm"] + _no_split_modules = ["IncontextAttentionBlock"] + _repeated_blocks = ["IncontextAttentionBlock"] + _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, + window_size: tuple = (-1, -1), + qk_norm: bool = True, + 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, + sparse_type: int = 0, + log_scale: float = 0.0, + ): + 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.window_size = window_size + self.qk_norm = qk_norm + 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 + self.sparse_type = sparse_type + self.log_scale = log_scale + + # [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.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( + [ + IncontextAttentionBlock( + dim, + ffn_dim, + num_heads, + window_size, + qk_norm, + cross_attn_norm, + eps, + refer_stride, + use_img_emb=use_img_emb, + sparse_type=sparse_type, + log_scale=log_scale, + ) + 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) + + # initialize weights + self.init_weights() + self.gradient_checkpointing = False + self.block_masks = {} + self.block_mask_grid_sizes = {} + + def create_mask(self, origin_len, origin_area, device): + origin_latent_f = origin_len // 4 + 1 + hw = int(np.prod(origin_area).item() // 256) + + 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, *args, method, **kwargs): + return getattr(self, method)(*args, **kwargs) + + def forward_ref( + self, + x_ref, + grid_sizes, + k_cache, + v_cache, + clip_fea_ref, + y_ref, + context_ref, + seq_len_ref, + t, + ): + device = self.patch_embedding.weight.device + # [reference] + x_ref = [torch.cat([u, v], dim=0) for u, v in zip(x_ref, y_ref)] + # embeddings + x_ref = [self.patch_embedding(u.unsqueeze(0)) for u in x_ref] + grid_sizes_ref = torch.stack([torch.tensor(u.shape[2:], dtype=torch.long) for u in x_ref]) + x_ref = [u.flatten(2).transpose(1, 2) for u in x_ref] + seq_lens_ref = torch.tensor([u.size(1) for u in x_ref], dtype=torch.long) + assert seq_lens_ref.max() <= seq_len_ref + x_ref = torch.cat([torch.cat([u, u.new_zeros(1, seq_len_ref - u.size(1), u.size(2))], dim=1) for u in x_ref]) + + assert (self.dim % self.num_heads) == 0 and (self.dim // self.num_heads) % 2 == 0 + d = self.dim // self.num_heads + + if self.refer_offset_t < 0: + self.refer_offset_t = grid_sizes[0][0].item() + if self.refer_offset_h < 0: + self.refer_offset_h = grid_sizes[0][1].item() + if self.refer_offset_w < 0: + self.refer_offset_w = 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) + + # time embeddings ref + with torch.amp.autocast(device_type="cuda", dtype=torch.float32): + e_ref = self.time_embedding(sinusoidal_embedding_1d(self.freq_dim, t * 0 + 1).float()) + e0_ref = self.time_projection(e_ref).unflatten(1, (6, self.dim)) + assert e_ref.dtype == torch.float32 and e0_ref.dtype == torch.float32 + + # [context_ref] + context_ref = self.text_embedding( + torch.stack([torch.cat([u, u.new_zeros(self.text_len - u.size(0), u.size(1))]) for u in context_ref]) + ) + + if self.use_img_emb: + context_clip_ref = self.img_emb(clip_fea_ref) + context_ref = torch.concat([context_clip_ref, context_ref], dim=1) + + context_lens = None + # arguments + kwargs = { + "e_ref": e0_ref, + "grid_sizes_ref": grid_sizes_ref, + "freqs_ref": self.freqs_ref, + "context_ref": context_ref, + "context_lens": context_lens, + } + + for idx, block in enumerate(self.blocks): + if torch.is_grad_enabled() and self.gradient_checkpointing: + x_ref = self._gradient_checkpointing_func( + block.forward_ref, + x_ref, + idx, + k_cache, + v_cache, + **kwargs, + ) + else: + x_ref = block(x_ref, idx, k_cache, v_cache, method="forward_ref", **kwargs) + + def forward_gen( + self, + x, + k_cache, + v_cache, + clip_fea, + y, + context, + seq_len, + t, + grid_sizes_ref, + origin_len, + origin_area, + is_uncondtion=False, + ): + # [denoising] + # params + device = self.patch_embedding.weight.device + x = [torch.cat([u, v], dim=0) for u, v in zip(x, y)] + # embeddings + x = [self.patch_embedding(u.unsqueeze(0)) for u in x] + grid_sizes = torch.stack([torch.tensor(u.shape[2:], dtype=torch.long) for u in x]) + x = [u.flatten(2).transpose(1, 2) for u in x] + seq_lens = torch.tensor([u.size(1) for u in x], dtype=torch.long) + assert seq_lens.max() <= seq_len + x = torch.cat([torch.cat([u, u.new_zeros(1, seq_len - u.size(1), u.size(2))], dim=1) for u in x]) + + assert (self.dim % self.num_heads) == 0 and (self.dim // self.num_heads) % 2 == 0 + d = self.dim // self.num_heads + 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) + + if self.refer_offset_t < 0: + self.refer_offset_t = grid_sizes[0][0].item() + if self.refer_offset_h < 0: + self.refer_offset_h = grid_sizes[0][1].item() + if self.refer_offset_w < 0: + self.refer_offset_w = 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) + + # time embeddings + with torch.amp.autocast(device_type="cuda", dtype=torch.float32): + e = self.time_embedding(sinusoidal_embedding_1d(self.freq_dim, t).float()) + e0 = self.time_projection(e).unflatten(1, (6, self.dim)) + assert e.dtype == torch.float32 and e0.dtype == torch.float32 + + # [context] + context_lens = None + context = self.text_embedding( + torch.stack([torch.cat([u, u.new_zeros(self.text_len - u.size(0), u.size(1))]) for u in context]) + ) + + if self.use_img_emb: + context_clip = self.img_emb(clip_fea) + context = torch.concat([context_clip, context], dim=1) + + block_mask_id = (origin_len, origin_area[0], origin_area[1]) + if block_mask_id not in self.block_masks: + self.block_masks[block_mask_id] = self.create_mask(origin_len, origin_area, x.device) + block_mask = self.block_masks[block_mask_id] + + # arguments + kwargs = { + "e": e0, + "block_mask": block_mask, + "grid_sizes": grid_sizes, + "freqs": self.freqs, + "context": context, + "grid_sizes_ref": grid_sizes_ref, + "freqs_ref": self.freqs_ref, + "context_lens": context_lens, + "origin_area": origin_area, + "origin_len": origin_len, + } + + for idx, block in enumerate(self.blocks): + if is_uncondtion and idx == 9: + continue + if torch.is_grad_enabled() and self.gradient_checkpointing: + x = self._gradient_checkpointing_func( + block.forward_gen, + x, + idx, + k_cache, + v_cache, + **kwargs, + ) + else: + x = block(x, idx, k_cache, v_cache, method="forward_gen", **kwargs) + + # head + x = self.head(x, e) + + # unpatchify + x = self.unpatchify(x, grid_sizes) + return [u.float() for u in x] + + 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 + + def init_weights(self): + # basic init + for m in self.modules(): + if isinstance(m, nn.Linear): + nn.init.xavier_uniform_(m.weight) + if m.bias is not None: + nn.init.zeros_(m.bias) + + # init embeddings + nn.init.xavier_uniform_(self.patch_embedding.weight.flatten(1)) + for m in self.text_embedding.modules(): + if isinstance(m, nn.Linear): + nn.init.normal_(m.weight, std=0.02) + for m in self.time_embedding.modules(): + if isinstance(m, nn.Linear): + nn.init.normal_(m.weight, std=0.02) + + # init output layer + nn.init.zeros_(self.head.head.weight) + + def load_from_official_state_dict(self, state_dict): + """Load weights from the official Wan-Animate-2 checkpoint.""" + self.load_state_dict(state_dict, strict=True) diff --git a/src/diffusers/modular_pipelines/wan/__init__.py b/src/diffusers/modular_pipelines/wan/__init__.py index 284b6c9fa436..0e0f06297311 100644 --- a/src/diffusers/modular_pipelines/wan/__init__.py +++ b/src/diffusers/modular_pipelines/wan/__init__.py @@ -21,11 +21,13 @@ _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"] = ["WanBlocks"] _import_structure["modular_blocks_wan22"] = ["Wan22Blocks"] _import_structure["modular_blocks_wan22_i2v"] = ["Wan22Image2VideoBlocks"] _import_structure["modular_blocks_wan_i2v"] = ["WanImage2VideoAutoBlocks"] _import_structure["modular_pipeline"] = [ + "WanAnimate2ModularPipeline", "Wan22Image2VideoModularPipeline", "Wan22ModularPipeline", "WanImage2VideoModularPipeline", @@ -42,10 +44,12 @@ from .modular_blocks_wan import WanBlocks from .modular_blocks_wan22 import Wan22Blocks from .modular_blocks_wan22_i2v import Wan22Image2VideoBlocks + from .modular_blocks_wan_animate_2 import WanAnimate2Blocks from .modular_blocks_wan_i2v import WanImage2VideoAutoBlocks from .modular_pipeline import ( Wan22Image2VideoModularPipeline, Wan22ModularPipeline, + WanAnimate2ModularPipeline, WanImage2VideoModularPipeline, WanModularPipeline, ) diff --git a/src/diffusers/modular_pipelines/wan/modular_blocks_wan_animate_2.py b/src/diffusers/modular_pipelines/wan/modular_blocks_wan_animate_2.py new file mode 100644 index 000000000000..21f8cd1bab46 --- /dev/null +++ b/src/diffusers/modular_pipelines/wan/modular_blocks_wan_animate_2.py @@ -0,0 +1,462 @@ +# 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 +import torch.nn.functional as F + +from ...configuration_utils import FrozenDict +from ...guiders import ClassifierFreeGuidance +from ...models import WanAnimate2Transformer3DModel +from ...schedulers import DPMSolverMultistepScheduler +from ...utils import logging +from ..modular_pipeline import ( + BlockState, + LoopSequentialPipelineBlocks, + ModularPipelineBlocks, + SequentialPipelineBlocks, +) +from ..modular_pipeline_utils import ComponentSpec, ConfigSpec, InputParam, OutputParam +from .decoders import WanVaeDecoderStep +from .encoders import WanTextEncoderStep +from .modular_pipeline import WanModularPipeline + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +def _get_sampling_sigmas(sampling_steps, shift): + sigma = np.linspace(1, 0, sampling_steps + 1)[:sampling_steps] + sigma = shift * sigma / (1 + (shift - 1) * sigma) + return sigma + + +class WanAnimate2ImageEncoderStep(ModularPipelineBlocks): + """Encode reference image with CLIP + VAE, and driving video with VAE.""" + + model_name = "wan" + + @property + def expected_components(self): + return [ + ComponentSpec("image_encoder", None), + ComponentSpec("vae", None), + ComponentSpec("image_processor", None), + ] + + @property + def description(self): + return "Encode reference image (CLIP + VAE) and driving video (VAE) for Wan-Animate-2." + + @property + def inputs(self): + return [ + InputParam("image", required=True, type_hint=object, description="Reference character image."), + InputParam("driving_video", required=True, type_hint=list, description="Driving video frames."), + InputParam("height", required=True, type_hint=int), + InputParam("width", required=True, type_hint=int), + InputParam("prompt_ref", required=False, type_hint=str, default="人物动作的参考视频"), + ] + + @property + def outputs(self): + return [ + OutputParam("clip_fea", type_hint=torch.Tensor, description="CLIP features of reference image."), + OutputParam("ref_latents", type_hint=torch.Tensor, description="VAE latents of reference image."), + OutputParam("condition_latents", type_hint=torch.Tensor, description="VAE latents of driving video."), + ] + + @torch.no_grad() + def __call__(self, components, state): + device = state.device + dtype = components.transformer.dtype + + # CLIP encode reference image + image = components.image_processor(images=state.image, return_tensors="pt").to(device) + image_embeds = components.image_encoder(**image, output_hidden_states=True) + state.clip_fea = image_embeds.hidden_states[-2].to(dtype) + + # VAE encode reference image + ref_pixels = components.image_processor(images=state.image, return_tensors="pt").to( + device=device, dtype=components.vae.dtype + ) + ref_latents = components.vae.encode(ref_pixels) + latents_mean = torch.tensor(components.vae.config.latents_mean).view(1, -1, 1, 1, 1).to(ref_latents) + latents_recip_std = 1.0 / torch.tensor(components.vae.config.latents_std).view(1, -1, 1, 1, 1).to(ref_latents) + state.ref_latents = (ref_latents - latents_mean) * latents_recip_std + + # VAE encode driving video + driving_pixels = state.driving_video.to(device=device, dtype=components.vae.dtype) + condition_latents = components.vae.encode(driving_pixels) + state.condition_latents = (condition_latents - latents_mean) * latents_recip_std + + return components, state + + +class WanAnimate2SetTimestepsStep(ModularPipelineBlocks): + """Set timesteps using custom sigma computation for flow matching.""" + + model_name = "wan" + + @property + def expected_components(self): + return [ComponentSpec("scheduler", DPMSolverMultistepScheduler)] + + @property + def description(self): + return "Set timesteps for Wan-Animate-2 with custom sigmas." + + @property + def inputs(self): + return [ + InputParam("num_inference_steps", required=True, type_hint=int), + InputParam("sample_shift", required=False, type_hint=float, default=5.0), + ] + + @property + def outputs(self): + return [OutputParam("timesteps", type_hint=torch.Tensor)] + + @torch.no_grad() + def __call__(self, components, state): + sigmas = _get_sampling_sigmas(state.num_inference_steps, state.sample_shift) + components.scheduler.set_timesteps(sigmas=sigmas, device=state.device) + state.timesteps = components.scheduler.timesteps + return components, state + + +class WanAnimate2PrepareLatentsStep(ModularPipelineBlocks): + """Prepare noise latents and encode reference (forward_ref -> KV cache).""" + + model_name = "wan" + + @property + def expected_components(self): + return [ComponentSpec("transformer", WanAnimate2Transformer3DModel)] + + @property + def description(self): + return "Prepare noise latents and encode reference video to cache KV." + + @property + def inputs(self): + return [ + InputParam("ref_latents", required=True, type_hint=torch.Tensor), + InputParam("clip_fea_ref", required=True, type_hint=torch.Tensor), + InputParam("condition_latents", required=True, type_hint=torch.Tensor), + InputParam("prompt_ref_embeds", required=True, type_hint=torch.Tensor), + InputParam("height", required=True, type_hint=int), + InputParam("width", required=True, type_hint=int), + InputParam("clip_len", required=True, type_hint=int), + InputParam("generator", required=False, type_hint=torch.Generator), + InputParam("timesteps", required=True, type_hint=torch.Tensor), + ] + + @property + def outputs(self): + return [ + OutputParam("latents", type_hint=torch.Tensor), + OutputParam("k_cache", type_hint=dict), + OutputParam("v_cache", type_hint=dict), + OutputParam("grid_sizes_ref", type_hint=torch.Tensor), + ] + + @torch.no_grad() + def __call__(self, components, state): + device = state.device + dtype = components.transformer.dtype + + latent_h = state.height // 8 + latent_w = state.width // 8 + clip_len = state.clip_len + lat_t = (clip_len + 3) // 4 + 1 + + # Prepare noise + noise = torch.randn( + 16, lat_t, latent_h, latent_w, device=device, dtype=torch.float32, generator=state.generator + ) + state.latents = [noise] + + # Prepare grid sizes for reference + ref_shape = state.condition_latents.shape[2:] + state.grid_sizes_ref = torch.tensor([ref_shape], dtype=torch.long) + + # KV cache + state.k_cache = {} + state.v_cache = {} + + # Reference encoding (forward_ref) + max_seq_len_ref = int(math.ceil(np.prod(ref_shape))) + + # Prepare y_ref (mask + ref_latents) + mask_ref = torch.zeros(1, 4, *state.ref_latents.shape[2:], device=device, dtype=dtype) + mask_ref[:, :, 0:1] = 1 + mask_ref = mask_ref.view(1, -1, 4, *state.ref_latents.shape[2:]).transpose(1, 2).squeeze(0) + y_ref = torch.cat([mask_ref, state.ref_latents[0]], dim=0) + + t_ref = torch.tensor([state.timesteps[0].item()], device=device, dtype=dtype) + + components.transformer( + [state.ref_latents[0]], + grid_sizes=state.grid_sizes_ref, + k_cache=state.k_cache, + v_cache=state.v_cache, + clip_fea_ref=state.clip_fea_ref, + y_ref=[y_ref], + context_ref=[state.prompt_ref_embeds[0]], + seq_len_ref=max_seq_len_ref, + t=t_ref, + method="forward_ref", + ) + + return components, state + + +class WanAnimate2LoopBeforeDenoiser(ModularPipelineBlocks): + """Prepare latent model input for the denoiser.""" + + model_name = "wan" + + @property + def description(self): + return "Prepare latent model input within the denoising loop." + + @property + def inputs(self): + return [InputParam("latents", required=True, type_hint=list)] + + @torch.no_grad() + def __call__(self, components, state, i, t): + state.latent_model_input = state.latents[0] + return components, state + + +class WanAnimate2LoopDenoiser(ModularPipelineBlocks): + """Denoiser that calls forward_gen with cached KV.""" + + model_name = "wan" + + def __init__(self, guider_input_fields=None): + if guider_input_fields is None: + guider_input_fields = {"context": ("prompt_embeds", "negative_prompt_embeds")} + self._guider_input_fields = guider_input_fields + super().__init__() + + @property + def expected_components(self): + return [ + ComponentSpec( + "guider", + ClassifierFreeGuidance, + config=FrozenDict({"guidance_scale": 3.0}), + default_creation_method="from_config", + ), + ComponentSpec("transformer", WanAnimate2Transformer3DModel), + ] + + @property + def description(self): + return "Denoiser step that calls forward_gen with cached KV for Wan-Animate-2." + + @property + def inputs(self): + inputs = [InputParam("num_inference_steps", required=True, type_hint=int)] + guider_names = [] + for v in self._guider_input_fields.values(): + if isinstance(v, tuple): + guider_names.extend(v) + else: + guider_names.append(v) + for name in guider_names: + inputs.append(InputParam(name=name, required=True, type_hint=torch.Tensor)) + return inputs + + @torch.no_grad() + def __call__(self, components, state, i, t): + components.guider.set_state(step=i, num_inference_steps=state.num_inference_steps, timestep=t) + guider_state = components.guider.prepare_inputs_from_block_state(state, self._guider_input_fields) + + for batch in guider_state: + components.guider.prepare_models(components.transformer) + cond_kwargs = batch.as_dict() + cond_kwargs = { + k: v.to(state.dtype) if isinstance(v, torch.Tensor) else v + for k, v in cond_kwargs.items() + if k in self._guider_input_fields + } + + is_uncond = batch.guidance_identifier == "pred_uncond" + batch.noise_pred = components.transformer( + state.latents, + k_cache=state.k_cache, + v_cache=state.v_cache, + clip_fea=state.clip_fea, + y=state.y, + seq_len=state.max_seq_len, + t=t.expand(1), + grid_sizes_ref=state.grid_sizes_ref, + origin_len=state.origin_len, + origin_area=state.origin_area, + method="forward_gen", + is_uncondtion=is_uncond, + **cond_kwargs, + ) + if isinstance(batch.noise_pred, list): + batch.noise_pred = batch.noise_pred[0] + components.guider.cleanup_models(components.transformer) + + state.noise_pred = components.guider(guider_state)[0] + return components, state + + +class WanAnimate2DenoiseLoopWrapper(LoopSequentialPipelineBlocks): + """Denoise loop for Wan-Animate-2: before_denoiser -> denoiser -> after_denoiser (scheduler step).""" + + model_name = "wan" + sub_blocks = [WanAnimate2LoopBeforeDenoiser, WanAnimate2LoopDenoiser] + + @property + def description(self): + return "Denoise loop for Wan-Animate-2 using forward_gen with cached KV." + + @property + def inputs(self): + return [ + InputParam("latents", required=True, type_hint=list), + InputParam("k_cache", required=True, type_hint=dict), + InputParam("v_cache", required=True, type_hint=dict), + InputParam("timesteps", required=True, type_hint=torch.Tensor), + ] + + @property + def outputs(self): + return [OutputParam("latents", type_hint=torch.Tensor)] + + @torch.no_grad() + def after_denoiser(self, components, state, i, t): + temp_x0 = components.scheduler.step( + state.noise_pred.unsqueeze(0), + t, + state.latents[0].unsqueeze(0), + return_dict=False, + )[0] + state.latents[0] = temp_x0.squeeze(0) + return components, state + + +# ==================== +# 1. CORE DENOISE +# ==================== + + +# auto_docstring +class WanAnimate2CoreDenoiseStep(SequentialPipelineBlocks): + """ + Core denoise block for Wan-Animate-2: set_timesteps -> prepare_latents (with ref encoding) -> denoise loop. + + Components: + transformer (`WanAnimate2Transformer3DModel`) scheduler (`DPMSolverMultistepScheduler`) guider + (`ClassifierFreeGuidance`) + + Inputs: + num_inference_steps (`int`): Number of denoising steps. + sample_shift (`float`): Shift for sigma computation. + ref_latents (`Tensor`): VAE latents of reference image. + condition_latents (`Tensor`): VAE latents of driving video. + clip_fea (`Tensor`): CLIP features of reference image. + clip_fea_ref (`Tensor`): CLIP features of driving video. + prompt_embeds (`Tensor`): Text embeddings. + negative_prompt_embeds (`Tensor`): Negative text embeddings. + prompt_ref_embeds (`Tensor`): Reference text embeddings. + height (`int`): Output height. + width (`int`): Output width. + clip_len (`int`): Frames per segment. + generator (`Generator`): Random generator. + + Outputs: + latents (`Tensor`): Denoised latents. + """ + + model_name = "wan" + block_classes = [ + WanAnimate2SetTimestepsStep, + WanAnimate2PrepareLatentsStep, + WanAnimate2DenoiseLoopWrapper, + ] + block_names = ["set_timesteps", "prepare_latents", "denoise"] + + @property + def description(self): + return "Core denoise block for Wan-Animate-2." + + @property + def outputs(self): + return [OutputParam.template("latents")] + + +# ==================== +# 2. FULL BLOCKS +# ==================== + + +# auto_docstring +class WanAnimate2Blocks(SequentialPipelineBlocks): + """ + Modular pipeline for character animation using Wan-Animate-2. + + Components: + text_encoder (`UMT5EncoderModel`) tokenizer (`AutoTokenizer`) image_encoder (`CLIPVisionModel`) + transformer (`WanAnimate2Transformer3DModel`) scheduler (`DPMSolverMultistepScheduler`) guider + (`ClassifierFreeGuidance`) vae (`AutoencoderKLWan`) video_processor (`VideoProcessor`) + + Inputs: + prompt (`str`): Text prompt describing the character. + negative_prompt (`str`): Negative prompt. + prompt_ref (`str`): Reference prompt for driving video. + image (`PIL.Image`): Reference character image. + driving_video (`list`): Driving video frames. + height (`int`): Output height. + width (`int`): Output width. + clip_len (`int`): Frames per segment. + num_inference_steps (`int`): Number of denoising steps. + sample_shift (`float`): Shift for sigma computation. + generator (`Generator`): Random generator. + output_type (`str`): Output format. + + Outputs: + videos (`list`): The generated videos. + """ + + model_name = "wan" + block_classes = [ + WanTextEncoderStep, + WanAnimate2ImageEncoderStep, + WanAnimate2CoreDenoiseStep, + WanVaeDecoderStep, + ] + block_names = [ + "text_encoder", + "image_encoder", + "denoise", + "decode", + ] + + @property + def description(self): + return "Modular pipeline for character animation using Wan-Animate-2." + + @property + def outputs(self): + return [OutputParam.template("videos")] diff --git a/src/diffusers/modular_pipelines/wan/modular_pipeline.py b/src/diffusers/modular_pipelines/wan/modular_pipeline.py index a360440c9251..74069843b714 100644 --- a/src/diffusers/modular_pipelines/wan/modular_pipeline.py +++ b/src/diffusers/modular_pipelines/wan/modular_pipeline.py @@ -139,3 +139,13 @@ class Wan22Image2VideoModularPipeline(Wan22ModularPipeline): """ default_blocks_name = "Wan22Image2VideoBlocks" + + +class WanAnimate2ModularPipeline(WanModularPipeline): + """ + 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" diff --git a/src/diffusers/pipelines/__init__.py b/src/diffusers/pipelines/__init__.py index 6c0c8667aab5..84862742a670 100644 --- a/src/diffusers/pipelines/__init__.py +++ b/src/diffusers/pipelines/__init__.py @@ -434,6 +434,7 @@ "WanVideoToVideoPipeline", "WanVACEPipeline", "WanAnimatePipeline", + "WanAnimate2Pipeline", ] _import_structure["kandinsky5"] = [ "Kandinsky5T2VPipeline", @@ -895,6 +896,7 @@ ) from .visualcloze import VisualClozeGenerationPipeline, VisualClozePipeline from .wan import ( + WanAnimate2Pipeline, WanAnimatePipeline, WanImageToVideoPipeline, WanPipeline, diff --git a/src/diffusers/pipelines/pipeline_utils.py b/src/diffusers/pipelines/pipeline_utils.py index a683973df5d7..60e16bdd47d8 100644 --- a/src/diffusers/pipelines/pipeline_utils.py +++ b/src/diffusers/pipelines/pipeline_utils.py @@ -30,22 +30,36 @@ import requests import torch from huggingface_hub import ( - DDUFEntry, ModelCard, create_repo, - get_cached_repo_tree, hf_hub_download, model_info, - read_dduf_file, snapshot_download, ) -from huggingface_hub.errors import CachedRepoTreeNotFoundError +try: + from huggingface_hub import DDUFEntry, read_dduf_file +except ImportError: + DDUFEntry = None + read_dduf_file = None +try: + from huggingface_hub import get_cached_repo_tree +except ImportError: + get_cached_repo_tree = None +try: + from huggingface_hub.errors import CachedRepoTreeNotFoundError +except ImportError: + class CachedRepoTreeNotFoundError(Exception): + pass from huggingface_hub.utils import ( HfHubHTTPError, LocalEntryNotFoundError, - OfflineModeIsEnabled, validate_hf_hub_args, ) +try: + from huggingface_hub.utils import OfflineModeIsEnabled +except ImportError: + class OfflineModeIsEnabled(Exception): + pass from packaging import version from tqdm.auto import tqdm from typing_extensions import Self diff --git a/src/diffusers/pipelines/wan/__init__.py b/src/diffusers/pipelines/wan/__init__.py index ad51a52f9242..3eac9b5a666b 100644 --- a/src/diffusers/pipelines/wan/__init__.py +++ b/src/diffusers/pipelines/wan/__init__.py @@ -24,6 +24,7 @@ else: _import_structure["pipeline_wan"] = ["WanPipeline"] _import_structure["pipeline_wan_animate"] = ["WanAnimatePipeline"] + _import_structure["pipeline_wan_animate_2"] = ["WanAnimate2Pipeline"] _import_structure["pipeline_wan_i2v"] = ["WanImageToVideoPipeline"] _import_structure["pipeline_wan_vace"] = ["WanVACEPipeline"] _import_structure["pipeline_wan_video2video"] = ["WanVideoToVideoPipeline"] @@ -37,6 +38,7 @@ else: from .pipeline_wan import WanPipeline from .pipeline_wan_animate import WanAnimatePipeline + from .pipeline_wan_animate_2 import WanAnimate2Pipeline from .pipeline_wan_i2v import WanImageToVideoPipeline from .pipeline_wan_vace import WanVACEPipeline from .pipeline_wan_video2video import WanVideoToVideoPipeline diff --git a/src/diffusers/pipelines/wan/pipeline_wan_animate_2.py b/src/diffusers/pipelines/wan/pipeline_wan_animate_2.py new file mode 100644 index 000000000000..ef1671441020 --- /dev/null +++ b/src/diffusers/pipelines/wan/pipeline_wan_animate_2.py @@ -0,0 +1,731 @@ +# 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 math +from typing import Callable + +import cv2 +import numpy as np +import torch +import torch.nn.functional as F + +from ...image_processor import PipelineImageInput +from ...loaders import WanLoraLoaderMixin +from ...models import AutoencoderKLWan, WanAnimate2Transformer3DModel +from ...schedulers import DPMSolverMultistepScheduler +from ...utils import logging +from ...video_processor import VideoProcessor +from ..pipeline_utils import DiffusionPipeline +from .pipeline_output import WanPipelineOutput + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +def get_sampling_sigmas(sampling_steps, shift): + sigma = np.linspace(1, 0, sampling_steps + 1)[:sampling_steps] + sigma = shift * sigma / (1 + (shift - 1) * sigma) + return sigma + + +def retrieve_timesteps(scheduler, num_inference_steps=None, device=None, sigmas=None, **kwargs): + if sigmas is not None: + accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) + if not accept_sigmas: + raise ValueError( + f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" + f" sigmas schedules." + ) + scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs) + timesteps = scheduler.timesteps + num_inference_steps = len(timesteps) + else: + scheduler.set_timesteps(num_inference_steps, device=device, **kwargs) + timesteps = scheduler.timesteps + return timesteps, num_inference_steps + + +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 + + +CLIP_MEAN = [0.48145466, 0.4578275, 0.40821073] +CLIP_STD = [0.26862954, 0.26130258, 0.27577711] + + +def get_frame_indices(frame_num, video_fps, clip_length, train_fps): + """Resample video frames to target fps.""" + times = np.arange(0, clip_length) / train_fps + frame_indices = np.round(times * video_fps).astype(int) + return np.clip(frame_indices, 0, frame_num - 1).tolist() + + +def padding_resize(img_ori, height, width, padding_color=(0, 0, 0), interpolation=cv2.INTER_LINEAR): + """Letterbox resize: keep aspect ratio + black padding to exact (height, width).""" + ori_h, ori_w = img_ori.shape[:2] + channel = img_ori.shape[2] if img_ori.ndim == 3 else 1 + img_pad = np.zeros((height, width, channel), dtype=np.uint8) + img_pad[:] = padding_color + + if ori_h / ori_w > height / width: + new_w = int(height / ori_h * ori_w) + img = cv2.resize(img_ori, (new_w, height), interpolation=interpolation) + padding = (width - new_w) // 2 + if img.ndim == 2: + img = img[:, :, np.newaxis] + img_pad[:, padding : padding + new_w, :] = img + return img_pad, {"padding_type": "width", "padding": padding, "side_long": new_w} + else: + new_h = int(width / ori_w * ori_h) + img = cv2.resize(img_ori, (width, new_h), interpolation=interpolation) + padding = (height - new_h) // 2 + if img.ndim == 2: + img = img[:, :, np.newaxis] + img_pad[padding : padding + new_h, :, :] = img + return img_pad, {"padding_type": "height", "padding": padding, "side_long": new_h} + + +def resize_by_area(image, target_area, divisor=16): + """Resize keeping aspect ratio targeting area, pad to exact dims. Returns (image, padding_info).""" + h, w = image.shape[:2] + aspect_ratio = w / h + new_h = math.sqrt(target_area / aspect_ratio) + new_w = target_area / new_h + new_w, new_h = int((new_w // divisor) * divisor), int((new_h // divisor) * divisor) + interpolation = cv2.INTER_AREA if (new_w * new_h < w * h) else cv2.INTER_LINEAR + return padding_resize(image, new_h, new_w, interpolation=interpolation) + + +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 + with torch.amp.autocast(device_type="cuda", dtype=dtype): + out = image_encoder(pixel_values=videos, output_hidden_states=True) + return out.hidden_states[-2] + + +class WanAnimate2Pipeline(DiffusionPipeline, WanLoraLoaderMixin): + r""" + Pipeline for character animation using Wan-Animate-2. + + This pipeline takes a reference character image and a driving video, and generates a video where the character + is animated following the motion in the driving video. The model uses an in-context attention mechanism with + KV cache: a reference video is first encoded to cache K/V tensors, then the generation forward uses the cached + K/V with a block mask for frame-level sparse in-context attention. + + Args: + tokenizer ([`AutoTokenizer`]): + Tokenizer for the umT5 text encoder. + text_encoder ([`UMT5EncoderModel`]): + The umT5 text encoder. + image_encoder ([`CLIPVisionModel`]): + CLIP vision model for encoding the reference image. + transformer ([`WanAnimate2Transformer3DModel`]): + The Wan-Animate-2 transformer model. + scheduler ([`DPMSolverMultistepScheduler`]): + A scheduler for flow matching. + vae ([`AutoencoderKLWan`]): + The Wan VAE model. + """ + + model_cpu_offload_seq = "text_encoder->image_encoder->transformer->vae" + _callback_tensor_inputs = ["latents"] + + def __init__( + self, + tokenizer, + text_encoder, + vae: AutoencoderKLWan, + scheduler: DPMSolverMultistepScheduler, + image_encoder, + transformer: WanAnimate2Transformer3DModel, + ): + super().__init__() + + self.register_modules( + vae=vae, + text_encoder=text_encoder, + tokenizer=tokenizer, + image_encoder=image_encoder, + transformer=transformer, + scheduler=scheduler, + ) + + self.vae_scale_factor_temporal = self.vae.config.scale_factor_temporal if getattr(self, "vae", None) else 4 + self.vae_scale_factor_spatial = self.vae.config.scale_factor_spatial if getattr(self, "vae", None) else 8 + self.video_processor = VideoProcessor(vae_scale_factor=self.vae_scale_factor_spatial) + + def _get_t5_prompt_embeds(self, prompt, device=None, dtype=None, max_sequence_length=512): + device = device or self._execution_device + dtype = dtype or self.text_encoder.dtype + + prompt = [prompt] if isinstance(prompt, str) else prompt + + text_inputs = self.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 = self.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 encode_image(self, image, device=None): + device = device or self._execution_device + from transformers import CLIPImageProcessor + + image_processor = CLIPImageProcessor() + processed = image_processor(images=image, return_tensors="pt").to(device) + image_embeds = self.image_encoder(**processed, output_hidden_states=True) + return image_embeds.hidden_states[-2] + + def _encode_vae(self, video, device, dtype): + """Encode video to latents using VAE, with standardization.""" + video = video.to(device=device, dtype=dtype) + latents = self.vae.encode(video) + if hasattr(latents, "latent_dist"): + latents = latents.latent_dist.mode() + elif hasattr(latents, "latents"): + latents = latents.latents + elif isinstance(latents, (list, tuple)): + latents = latents[0] if isinstance(latents[0], torch.Tensor) else torch.stack(latents) + latents_mean = ( + torch.tensor(self.vae.config.latents_mean) + .view(1, self.vae.config.z_dim, 1, 1, 1) + .to(latents.device, latents.dtype) + ) + latents_recip_std = 1.0 / torch.tensor(self.vae.config.latents_std).view(1, self.vae.config.z_dim, 1, 1, 1).to( + latents.device, latents.dtype + ) + latents = (latents - latents_mean) * latents_recip_std + return latents + + def _decode_vae(self, latents, device): + """Decode latents to video using VAE, with destandardization.""" + latents = latents.to(self.vae.dtype) + latents_mean = ( + torch.tensor(self.vae.config.latents_mean) + .view(1, self.vae.config.z_dim, 1, 1, 1) + .to(latents.device, latents.dtype) + ) + latents_recip_std = 1.0 / torch.tensor(self.vae.config.latents_std).view(1, self.vae.config.z_dim, 1, 1, 1).to( + latents.device, latents.dtype + ) + latents = latents / latents_recip_std + latents_mean + out_frames = self.vae.decode(latents, return_dict=False)[0] + return out_frames + + def check_inputs(self, image, driving_video, prompt, height, width): + if image is None: + raise ValueError("Provide `image`. Cannot leave `image` undefined.") + if driving_video is None: + raise ValueError("Provide `driving_video`. Cannot leave `driving_video` undefined.") + if height % 16 != 0 or width % 16 != 0: + raise ValueError(f"`height` and `width` have to be divisible by 16 but are {height} and {width}.") + + @property + def guidance_scale(self): + return self._guidance_scale + + @property + def do_classifier_free_guidance(self): + return self._guidance_scale > 1 + + @property + def num_timesteps(self): + return self._num_timesteps + + @torch.no_grad() + def __call__( + self, + image: PipelineImageInput, + driving_video: list, + prompt: str | list[str] = None, + negative_prompt: str | list[str] = None, + prompt_ref: str = "人物动作的参考视频", + height: int = 800, + width: int = 640, + clip_len: int = 81, + first_num: int = 1, + fps: int = 24, + num_inference_steps: int = 40, + guidance_scale: float = 3.0, + sample_shift: float = 5.0, + flow_solver: str = "dpm", + seed: int = -1, + generator: torch.Generator | list[torch.Generator] | None = None, + output_type: str | None = "np", + return_dict: bool = True, + callback_on_step_end: Callable | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], + max_sequence_length: int = 512, + ): + r""" + The call function for character animation generation. + + Args: + image (`PipelineImageInput`): + The reference character image. + driving_video (`list`): + The driving video (list of PIL images or tensors) that provides motion. + prompt (`str` or `list[str]`): + The text prompt describing the character appearance and background. + negative_prompt (`str` or `list[str]`, *optional*): + The negative prompt for classifier-free guidance. + prompt_ref (`str`, defaults to `"人物动作的参考视频"`): + The reference prompt for the driving video context. + height (`int`, defaults to `800`): + The height of the generated video. + width (`int`, defaults to `640`): + The width of the generated video. + clip_len (`int`, defaults to `81`): + The number of frames in each inference segment. + first_num (`int`, defaults to `1`): + The number of conditioning frames from the previous segment. + fps (`int`, defaults to `24`): + The output video FPS. + num_inference_steps (`int`, defaults to `40`): + The number of denoising steps. + guidance_scale (`float`, defaults to `3.0`): + Guidance scale for classifier-free guidance. + sample_shift (`float`, defaults to `5.0`): + The shift parameter for sigma computation. + seed (`int`, defaults to `-1`): + Random seed. -1 means random. + output_type (`str`, defaults to `"np"`): + The output format. + return_dict (`bool`, defaults to `True`): + Whether to return a `WanPipelineOutput`. + """ + # 1. Check inputs + self.check_inputs(image, driving_video, prompt, height, width) + + self._guidance_scale = guidance_scale + device = self._execution_device + + if seed >= 0: + torch.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + np.random.seed(seed) + + if generator is None: + generator = torch.Generator(device=device) + if seed >= 0: + generator.manual_seed(seed) + + # 2. Preprocess reference image (letterbox resize — do this first to get actual dims) + ref_np = np.array(image) # PIL → numpy [H, W, C] + ref_pad, ref_padding_info = resize_by_area(ref_np, width * height, divisor=16) + actual_h, actual_w = ref_pad.shape[:2] + + # 3. Prepare driving video frames (FPS resampling + letterbox to match ref dims) + import decord + + vr = decord.VideoReader(driving_video if isinstance(driving_video, str) else None) + video_fps = vr.get_avg_fps() + frame_num = len(vr) + target_num = int(frame_num / video_fps * fps) + idxs = get_frame_indices(frame_num, video_fps, target_num, fps) + frames_np = vr.get_batch(idxs).asnumpy() # [T, H, W, C] uint8 + + cond_images_np = [] + for frame in frames_np: + img_pad, _ = padding_resize(frame, actual_h, actual_w) + cond_images_np.append(img_pad) + + driving_video = torch.tensor(np.stack(cond_images_np), dtype=torch.float32) # [T, H, W, C] + driving_video = driving_video / 127.5 - 1.0 # [-1, 1] + driving_video = driving_video.permute(3, 0, 1, 2).unsqueeze(0) # [1, C, T, H, W] + driving_video = driving_video.to(device, dtype=torch.float32) + + # Pad driving video to be a multiple of (clip_len - first_num) + real_frame_len = driving_video.shape[2] + effective_segment = clip_len - first_num + last_segment_frames = (real_frame_len - first_num) % effective_segment if real_frame_len > first_num else 0 + if last_segment_frames > 0: + num_padding = effective_segment - last_segment_frames + else: + num_padding = 0 + target_num_frames = real_frame_len + num_padding + + # Pad driving video using zigzag (reflect) strategy + 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) + + # 4. Encode prompt + prompt_embeds = self._get_t5_prompt_embeds(prompt, device=device, max_sequence_length=max_sequence_length) + negative_prompt_embeds = None + if self.do_classifier_free_guidance: + negative_prompt = negative_prompt or "" + negative_prompt_embeds = self._get_t5_prompt_embeds( + negative_prompt, device=device, max_sequence_length=max_sequence_length + ) + + # Reference prompt + prompt_ref_embeds = self._get_t5_prompt_embeds( + prompt_ref, device=device, max_sequence_length=max_sequence_length + ) + + # 5. Encode reference image (VAE + CLIP) + ref_tensor = torch.tensor(ref_pad, dtype=torch.float32) / 127.5 - 1.0 # [-1, 1] + image_pixels = ref_tensor.permute(2, 0, 1).unsqueeze(0).unsqueeze(2).to(device, dtype=torch.float32) + + # CLIP features from reference image (direct bicubic to 224×224 from tensor) + clip_fea = clip_visual_encode(self.image_encoder, ref_tensor.permute(2, 0, 1).to(device), device, self.transformer.dtype) + + # VAE encode reference image + ref_pixels = image_pixels.to(self.vae.dtype) + if ref_pixels.ndim == 4: + ref_pixels = ref_pixels.unsqueeze(2) # [B, C, H, W] -> [B, C, 1, H, W] + ref_latents = self.vae.encode(ref_pixels) + if hasattr(ref_latents, "latent_dist"): + ref_latents = ref_latents.latent_dist.mode() + elif hasattr(ref_latents, "latents"): + ref_latents = ref_latents.latents + elif isinstance(ref_latents, (list, tuple)): + ref_latents = torch.stack(ref_latents) if not isinstance(ref_latents[0], torch.Tensor) else ref_latents[0] + latents_mean = ( + torch.tensor(self.vae.config.latents_mean) + .view(1, self.vae.config.z_dim, 1, 1, 1) + .to(ref_latents.device, ref_latents.dtype) + ) + latents_recip_std = 1.0 / torch.tensor(self.vae.config.latents_std).view(1, self.vae.config.z_dim, 1, 1, 1).to( + ref_latents.device, ref_latents.dtype + ) + ref_latents = (ref_latents - latents_mean) * latents_recip_std + + # Derive latent dims from ACTUAL image size after resize_by_area (not requested height/width) + actual_h, actual_w = ref_pad.shape[:2] + latent_h = actual_h // self.vae_scale_factor_spatial + latent_w = actual_w // self.vae_scale_factor_spatial + + # Prepare reference i2v mask and y_ref + mask_ref = get_i2v_mask(1, latent_h, latent_w, 1, device=device).to(self.transformer.dtype) + ref_lat_0 = ref_latents[0] if ref_latents.ndim == 5 else ref_latents + y_ref = torch.cat([mask_ref, ref_lat_0], dim=0) + + # CLIP context for reference + clip_context = clip_fea + + # 5. Set up scheduler + if flow_solver == "euler": + from diffusers import FlowMatchEulerDiscreteScheduler + + sample_scheduler = FlowMatchEulerDiscreteScheduler( + num_train_timesteps=1000, + shift=sample_shift, + use_dynamic_shifting=False, + ) + else: + sample_scheduler = DPMSolverMultistepScheduler.from_config( + self.scheduler.config, + num_train_timesteps=1000, + flow_shift=sample_shift, + use_dynamic_shifting=False, + prediction_type="flow_prediction", + ) + sample_scheduler.set_timesteps(num_inference_steps, device=device) + timesteps = sample_scheduler.timesteps + + self._num_timesteps = len(timesteps) + + # 6. Segment-based generation loop + start = 0 + end = clip_len + all_out_frames = [] + out_frames = None + + num_segments = (target_num_frames - first_num + effective_segment - 1) // effective_segment + + for seg_idx in range(num_segments): + if start + first_num >= target_num_frames: + break + + mask_reft_len = first_num if start > 0 else 0 + + if target_num_frames - start < clip_len: + clip_len_actual = target_num_frames - start + else: + clip_len_actual = clip_len + + # VAE encode the driving video segment + cond_pixels = driving_video[:, :, start : start + clip_len_actual].to(self.vae.dtype) + condition_latents = self.vae.encode(cond_pixels) + if hasattr(condition_latents, "latent_dist"): + condition_latents = condition_latents.latent_dist.mode() + elif hasattr(condition_latents, "latents"): + condition_latents = condition_latents.latents + elif isinstance(condition_latents, (list, tuple)): + condition_latents = condition_latents[0] if isinstance(condition_latents[0], torch.Tensor) else torch.stack(condition_latents) + condition_latents = (condition_latents - latents_mean) * latents_recip_std + + # CLIP features from driving video first frame (direct bicubic to 224×224 from tensor) + condition_img = driving_video[0, :, 0] # [C, H, W] in [-1, 1] + condition_clip_context = clip_visual_encode( + self.image_encoder, condition_img, device, self.transformer.dtype + ) + + # Prepare condition y (mask + latents) + T = clip_len_actual + 1 + + # Encode condition y + if mask_reft_len > 0: + prev_frames = out_frames[0, :, -mask_reft_len:].clone().detach() + prev_frames_interp = F.interpolate( + prev_frames.permute(1, 0, 2, 3), size=(actual_h, actual_w), mode="bicubic" + ).permute(1, 0, 2, 3) + cond_y_input = torch.cat( + [prev_frames_interp, torch.zeros(3, T - mask_reft_len - 1, actual_h, actual_w, device=device)], + dim=1, + ).to(self.vae.dtype) + else: + cond_y_input = torch.zeros(3, T - 1, actual_h, actual_w, device=device).to(self.vae.dtype) + + y_reft = self.vae.encode(cond_y_input.unsqueeze(0)) + if hasattr(y_reft, "latent_dist"): + y_reft = y_reft.latent_dist.mode() + elif hasattr(y_reft, "latents"): + y_reft = y_reft.latents + elif isinstance(y_reft, (list, tuple)): + y_reft = y_reft[0] + y_reft = (y_reft - latents_mean) * latents_recip_std + if y_reft.ndim == 5: + y_reft = y_reft.squeeze(0) # [1, 16, T, H, W] -> [16, T, H, W] + + # Derive lat_t from actual VAE output shape + lat_t_y = y_reft.shape[1] # temporal dimension of y_reft latents + lat_t_cond = condition_latents.shape[2] if condition_latents.ndim == 5 else condition_latents.shape[1] + + msk_reft = get_i2v_mask(lat_t_y, latent_h, latent_w, mask_reft_len, device=device).to( + self.transformer.dtype + ) + y_reft = torch.cat([msk_reft, y_reft], dim=0) + + # Condition mask and latents + condition_msk_y = get_i2v_mask(lat_t_cond, latent_h, latent_w, clip_len_actual, device=device).to( + self.transformer.dtype + ) + cond_lat_0 = condition_latents[0] if condition_latents.ndim == 5 else condition_latents + condition_y = torch.cat([condition_msk_y, cond_lat_0], dim=0) + + y = torch.cat([y_ref, y_reft], dim=1) + + # Prepare grid sizes — use post-patch spatial dims (VAE 8x + patch 2x = 16x total) + if condition_latents.ndim == 5: + ref_shape = list(condition_latents.shape[2:]) # [T, H, W] pre-patch + else: + ref_shape = list(condition_latents.shape[1:]) + # After patch_embedding (1,2,2): spatial dims halved + ref_shape_post = [ref_shape[0], ref_shape[1] // 2, ref_shape[2] // 2] + grid_sizes_ref = torch.tensor([ref_shape_post], dtype=torch.long) + + # Noise latents temporal dim = y_ref(1) + y_reft/condition_y(T) = total y temporal dim + lat_t_noise = y.shape[1] if y.ndim == 4 else y.shape[2] + noise = torch.randn( + 16, + lat_t_noise, + latent_h, + latent_w, + dtype=torch.float32, + device=device, + generator=generator, + ) + + latents = [noise] + + # Prepare arguments for transformer + max_seq_len = int(math.ceil(np.prod([lat_t_noise, latent_h // 2, latent_w // 2]))) + max_seq_len_ref = int(math.ceil(np.prod(ref_shape) // 4)) if ref_shape else max_seq_len + + arg_c = { + "context": [prompt_embeds[0]], + "seq_len": max_seq_len, + "clip_fea": clip_context, + "y": [y], + "origin_len": clip_len_actual, + "origin_area": [actual_h, actual_w], + } + + arg_ref_c = { + "context_ref": [prompt_ref_embeds[0]], + "seq_len_ref": max_seq_len_ref, + "clip_fea_ref": condition_clip_context, + "y_ref": [condition_y], + } + + arg_null = None + if self.do_classifier_free_guidance: + arg_null = { + "context": [negative_prompt_embeds[0]], + "seq_len": max_seq_len, + "clip_fea": clip_context, + "y": [y], + "origin_len": clip_len_actual, + "origin_area": [actual_h, actual_w], + "is_uncondtion": True, + } + + # KV cache + k_cache = {} + v_cache = {} + + # Phase 1: encode reference — cast all inputs to transformer dtype + t_ref = torch.tensor([timesteps[0].item()], device=device, dtype=self.transformer.dtype) + with torch.amp.autocast(device_type="cuda", dtype=torch.bfloat16): + self.transformer( + [condition_latents[0].to(self.transformer.dtype)] if condition_latents.ndim == 5 else [condition_latents.to(self.transformer.dtype)], + grid_sizes=grid_sizes_ref, + k_cache=k_cache, + v_cache=v_cache, + clip_fea_ref=arg_ref_c["clip_fea_ref"].to(self.transformer.dtype), + y_ref=[y.to(self.transformer.dtype) for y in arg_ref_c["y_ref"]], + context_ref=[c.to(self.transformer.dtype) for c in arg_ref_c["context_ref"]], + seq_len_ref=max_seq_len_ref, + t=t_ref, + method="forward_ref", + ) + + # Phase 2: denoising loop + from tqdm import tqdm + + for i, t in tqdm(enumerate(timesteps), total=len(timesteps), desc=f"Segment {seg_idx+1}/{num_segments}"): + timestep = torch.stack([t]) + + with torch.amp.autocast(device_type="cuda", dtype=torch.bfloat16): + # Conditional + noise_pred_cond = self.transformer( + latents, + k_cache=k_cache, + v_cache=v_cache, + clip_fea=arg_c["clip_fea"], + y=arg_c["y"], + context=arg_c["context"], + seq_len=max_seq_len, + t=timestep, + grid_sizes_ref=grid_sizes_ref, + origin_len=arg_c["origin_len"], + origin_area=arg_c["origin_area"], + method="forward_gen", + ) + if isinstance(noise_pred_cond, list): + noise_pred_cond = noise_pred_cond[0] + + if self.do_classifier_free_guidance: + noise_pred_uncond = self.transformer( + latents, + k_cache=k_cache, + v_cache=v_cache, + clip_fea=arg_null["clip_fea"], + y=arg_null["y"], + context=arg_null["context"], + seq_len=max_seq_len, + t=timestep, + grid_sizes_ref=grid_sizes_ref, + origin_len=arg_null["origin_len"], + origin_area=arg_null["origin_area"], + method="forward_gen", + is_uncondtion=True, + ) + if isinstance(noise_pred_uncond, list): + noise_pred_uncond = noise_pred_uncond[0] + + noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_cond - noise_pred_uncond) + else: + noise_pred = noise_pred_cond + + # Scheduler step + temp_x0 = sample_scheduler.step( + noise_pred.unsqueeze(0), + t, + latents[0].unsqueeze(0), + return_dict=False, + generator=generator, + )[0] + latents[0] = temp_x0.squeeze(0) + + if callback_on_step_end is not None: + callback_kwargs = {} + for k in callback_on_step_end_tensor_inputs: + callback_kwargs[k] = locals()[k] + callback_outputs = callback_on_step_end(self, i, t, callback_kwargs) + latents[0] = ( + callback_outputs.pop("latents", latents)[0] + if isinstance(callback_outputs.get("latents"), list) + else latents[0] + ) + + # Decode + x0 = [latents[0].to(dtype=torch.float32)] + out_frames = self._decode_vae(x0[0][:, 1:], device) + + if start > 0: + out_frames = out_frames[:, :, mask_reft_len:] + + all_out_frames.append(out_frames) + start += effective_segment + end += effective_segment + + # Reset scheduler for next segment + sample_scheduler.set_timesteps(num_inference_steps, device=device) + timesteps = sample_scheduler.timesteps + + # Concatenate all segments + video = torch.cat(all_out_frames, dim=2)[:, :, :real_frame_len] + + # Remove letterbox padding (crop black borders) + p_info = ref_padding_info + if p_info["padding_type"] == "width": + video = video[:, :, :, :, p_info["padding"] : p_info["padding"] + p_info["side_long"]] + else: + video = video[:, :, :, p_info["padding"] : p_info["padding"] + p_info["side_long"], :] + + video = self.video_processor.postprocess_video(video, output_type=output_type) + + self.maybe_free_model_hooks() + + if not return_dict: + return (video,) + + return WanPipelineOutput(frames=video) diff --git a/src/diffusers/utils/dummy_pt_objects.py b/src/diffusers/utils/dummy_pt_objects.py index 9035efb3e6e2..aec1d980ca08 100644 --- a/src/diffusers/utils/dummy_pt_objects.py +++ b/src/diffusers/utils/dummy_pt_objects.py @@ -2250,6 +2250,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 417aa9ad18ff..60e3398abeee 100644 --- a/src/diffusers/utils/dummy_torch_and_transformers_objects.py +++ b/src/diffusers/utils/dummy_torch_and_transformers_objects.py @@ -4982,6 +4982,21 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch", "transformers"]) +class WanAnimate2Pipeline(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 WanAnimatePipeline(metaclass=DummyObject): _backends = ["torch", "transformers"] From a54425a5bda8896dbb18305f63966385fc020b44 Mon Sep 17 00:00:00 2001 From: "yiyi@huggingface.co" Date: Fri, 7 Aug 2026 04:33:41 +0000 Subject: [PATCH 02/19] Refactor Wan-Animate-2 to diffusers conventions Model (`transformer_wan_animate_2.py`): - Replace the `forward(*args, method=...)` dispatch and the split `forward_ref`/`forward_gen` with a single documented `forward(..., kv_cache_mode="extract"|"cached")` returning `Transformer2DModelOutput`, following the Flux2 KV-cache precedent. The `SelfAttention`/`CrossAttention` pre/post split becomes a regular `WanAnimate2Attention` (`AttentionModuleMixin`) with processors that run through `dispatch_attention_fn` - native SDPA by default, any backend via `set_attention_backend`; only the in-context generation path is pinned to `flex`, since its attention pattern is expressed as a `BlockMask`. The hard `flash_attn` requirement is gone. - `IncontextAttentionBlock` was a pure pass-through around `AttentionBlock`; merged into one `WanAnimate2TransformerBlock` (checkpoint keys lose the `.block.` segment, handled in the single-file mapping). - KV cache is a `WanAnimate2KVCache` object instead of bare dicts passed through `forward`. Accelerate hooks copy dict arguments, so the dict version breaks under `enable_model_cpu_offload` (the reference pass fills a copy and the generation pass KeyErrors); the object passes through by reference, and `_skip_keys = ["kv_cache"]` covers group offloading. - Remove all autocast in favour of the `transformer_wan.py` dtype discipline (fp32 modulation with `.type_as` casts at block boundaries), so the model runs natively in bf16 and is no longer CUDA-only in principle. Replace the local float64 `sinusoidal_embedding_1d` with the existing `Timesteps` class. - Remove dead code: the unreachable padding mask in the reference path (the pipeline always fills `seq_len` exactly), `init_weights`, `load_from_official_state_dict`, and the unused `window_size`/`qk_norm`/ `sparse_type`/`log_scale` config flags (`log_scale` ships as 0.0, making the flex `score_mod` a no-op). Pipeline: call sites updated to the merged forward, autocast wrappers replaced with explicit casts at the call boundary. Also: - Fill in `convert_wan_animate_2_transformer_to_diffusers` with the actual key mapping (block unwrap + attention renames); it was a prefix-strip no-op. - Revert the `pipeline_utils.py` try/except import shims - the stub exception classes silently break real `except OfflineModeIsEnabled` handling; upgrade `huggingface_hub` to the `setup.py` pin instead. - Drop the modular pipeline for now; it needs its own pass and is not part of the initial release surface. Numerics: the refactored model matches the reference implementation at 2.47e-05 max relative difference in fp32 over all 40 layers on real weights, and end-to-end outputs match the reference pipeline to a max pixel difference of 7e-5 (PSNR 119 dB) when kernels and environment are held fixed. The checkpoint key rename is a pure rename - all 1303 tensors bitwise identical. Co-Authored-By: Claude Opus 5 --- src/diffusers/loaders/single_file_utils.py | 31 +- .../transformers/transformer_wan_animate_2.py | 1272 +++++++---------- .../modular_pipelines/wan/__init__.py | 4 - .../wan/modular_blocks_wan_animate_2.py | 462 ------ .../modular_pipelines/wan/modular_pipeline.py | 10 - src/diffusers/pipelines/pipeline_utils.py | 24 +- .../pipelines/wan/pipeline_wan_animate_2.py | 105 +- 7 files changed, 620 insertions(+), 1288 deletions(-) delete mode 100644 src/diffusers/modular_pipelines/wan/modular_blocks_wan_animate_2.py diff --git a/src/diffusers/loaders/single_file_utils.py b/src/diffusers/loaders/single_file_utils.py index c22ddb9a3a18..b5c6846ba17c 100644 --- a/src/diffusers/loaders/single_file_utils.py +++ b/src/diffusers/loaders/single_file_utils.py @@ -3294,18 +3294,29 @@ def convert_wan_animate_2_transformer_to_diffusers(checkpoint, **kwargs): Converts the state dict of the Wan-Animate-2 transformer from the official checkpoint format to the diffusers format. """ - converted_state_dict = {} - - # Strip model.diffusion_model prefix if present - keys = list(checkpoint.keys()) - for k in keys: - if "model.diffusion_model." in k: - checkpoint[k.replace("model.diffusion_model.", "")] = checkpoint.pop(k) + 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.", + } - # The official checkpoint already uses the same key format as the diffusers model - # (blocks.N.block.*), so no remapping is needed. + converted_state_dict = {} for key in list(checkpoint.keys()): - converted_state_dict[key] = checkpoint.pop(key) + 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 diff --git a/src/diffusers/models/transformers/transformer_wan_animate_2.py b/src/diffusers/models/transformers/transformer_wan_animate_2.py index e9edd44b1f9c..36309a0001fc 100644 --- a/src/diffusers/models/transformers/transformer_wan_animate_2.py +++ b/src/diffusers/models/transformers/transformer_wan_animate_2.py @@ -13,228 +13,21 @@ # limitations under the License. import math -from functools import lru_cache, partial -import numpy as np import torch import torch.nn as nn -from torch.nn.attention.flex_attention import create_block_mask +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 -try: - from flash_attn_interface import flash_attn_varlen_func - - FLASH_VER = 3 -except ModuleNotFoundError: - try: - from flash_attn import flash_attn_varlen_func - - FLASH_VER = 2 - except ModuleNotFoundError: - flash_attn_varlen_func = None - FLASH_VER = None - -from torch.nn.attention.flex_attention import flex_attention as _flex_attention_raw - -# Lazy compile: compile on first call instead of at import time -_flex_compiled = None - - -def _get_compiled_flex_attention(): - global _flex_compiled - if _flex_compiled is None: - _flex_compiled = torch.compile(_flex_attention_raw, dynamic=False, mode="max-autotune", fullgraph=True) - return _flex_compiled - - -def flash_attention( - q, - k, - v, - q_lens=None, - k_lens=None, - dropout_p=0.0, - softmax_scale=None, - q_scale=None, - causal=False, - window_size=(-1, -1), - deterministic=False, - dtype=torch.bfloat16, -): - """ - q: [B, Lq, Nq, C1]. - k: [B, Lk, Nk, C1]. - v: [B, Lk, Nk, C2]. Nq must be divisible by Nk. - q_lens: [B]. - k_lens: [B]. - dropout_p: float. Dropout probability. - softmax_scale: float. The scaling of QK^T before applying softmax. - causal: bool. Whether to apply causal attention mask. - window_size: (left right). If not (-1, -1), apply sliding window local attention. - deterministic: bool. If True, slightly slower and uses more memory. - dtype: torch.dtype. Apply when dtype of q/k/v is not float16/bfloat16. - """ - half_dtypes = (torch.float16, torch.bfloat16) - assert dtype in half_dtypes - assert q.device.type == "cuda" and q.size(-1) <= 256 - - # params - b, lq, lk, out_dtype = q.size(0), q.size(1), k.size(1), q.dtype - - def half(x): - return x if x.dtype in half_dtypes else x.to(dtype) - - # preprocess query - if q_lens is None: - q = half(q.flatten(0, 1)) - q_lens = torch.tensor([lq] * b, dtype=torch.int32).to(device=q.device, non_blocking=True) - else: - q = half(torch.cat([u[:v] for u, v in zip(q, q_lens)])) - - # preprocess key, value - if k_lens is None: - k = half(k.flatten(0, 1)) - v = half(v.flatten(0, 1)) - k_lens = torch.tensor([lk] * b, dtype=torch.int32).to(device=k.device, non_blocking=True) - else: - k = half(torch.cat([u[:v] for u, v in zip(k, k_lens)])) - v = half(torch.cat([u[:v] for u, v in zip(v, k_lens)])) - - q = q.to(v.dtype) - k = k.to(v.dtype) - - if q_scale is not None: - q = q * q_scale - # apply attention - if FLASH_VER == 3: - # Note: dropout_p, window_size are not supported in FA3 now. - x = flash_attn_varlen_func( - q=q, - k=k, - v=v, - cu_seqlens_q=torch.cat([q_lens.new_zeros([1]), q_lens]) - .cumsum(0, dtype=torch.int32) - .to(q.device, non_blocking=True), - cu_seqlens_k=torch.cat([k_lens.new_zeros([1]), k_lens]) - .cumsum(0, dtype=torch.int32) - .to(q.device, non_blocking=True), - max_seqlen_q=lq, - max_seqlen_k=lk, - softmax_scale=softmax_scale, - causal=causal, - deterministic=deterministic, - )[0].unflatten(0, (b, lq)) - else: - assert FLASH_VER == 2 - x = flash_attn_varlen_func( - q=q, - k=k, - v=v, - cu_seqlens_q=torch.cat([q_lens.new_zeros([1]), q_lens]) - .cumsum(0, dtype=torch.int32) - .to(q.device, non_blocking=True), - cu_seqlens_k=torch.cat([k_lens.new_zeros([1]), k_lens]) - .cumsum(0, dtype=torch.int32) - .to(q.device, non_blocking=True), - max_seqlen_q=lq, - max_seqlen_k=lk, - dropout_p=dropout_p, - softmax_scale=softmax_scale, - causal=causal, - window_size=window_size, - deterministic=deterministic, - ).unflatten(0, (b, lq)) - - # output - return x.type(out_dtype) - - -def flex_attention( - q, - k, - v, - q_lens=None, - k_lens=None, - block_mask=None, - kernel_options=None, - dtype=torch.bfloat16, - score_mod=None, -): - """ - q: [B, Lq, Nq, C1]. - k: [B, Lk, Nk, C1]. - v: [B, Lk, Nk, C2]. Nq must be divisible by Nk. - q_lens: [B]. - k_lens: [B]. - dtype: torch.dtype. Apply when dtype of q/k/v is not float16/bfloat16. - """ - half_dtypes = (torch.float16, torch.bfloat16) - assert dtype in half_dtypes - assert q.device.type == "cuda" - lq, lk, out_dtype = q.size(1), k.size(1), q.dtype - - def half(x): - return x if x.dtype in half_dtypes else x.to(dtype) - - assert lq % 128 == 0, "q_len must be divisible by 128." - assert lk % 128 == 0, "k_len must be divisible by 128." - - # preprocess query - if q_lens is None: - q = half(q) - else: - q = half(q) - assert q_lens.max() == q_lens.min(), "varlen of query is not supported" - - # preprocess key, value - if k_lens is None: - k, v = half(k), half(v) - else: - k, v = half(k), half(v) - assert k_lens.max() == k_lens.min(), "varlen of key is not supported" - - q = q.to(v.dtype) - k = k.to(v.dtype) - - x = _get_compiled_flex_attention()( - query=q.transpose(2, 1), - key=k.transpose(2, 1), - value=v.transpose(2, 1), - block_mask=block_mask, - kernel_options=kernel_options, - score_mod=score_mod, - ).transpose(2, 1) - - return x.type(out_dtype) - - -def _score_mod_impl(score, b_idx, h_idx, q_idx, kv_idx, hw: int, log_scale: float): - condition = (kv_idx >= hw) & (kv_idx < 2 * hw) - return torch.where(condition, score + log_scale, score) - - -@lru_cache(maxsize=32) -def _get_score_mod(hw: int, log_scale: float = -1.0): - return partial(_score_mod_impl, hw=hw, log_scale=log_scale) - - -def sinusoidal_embedding_1d(dim, position): - # preprocess - assert dim % 2 == 0 - half = dim // 2 - position = position.type(torch.float64) - - # calculation - sinusoid = torch.outer(position, torch.pow(10000, -torch.arange(half).to(position).div(half))) - x = torch.cat([torch.cos(sinusoid), torch.sin(sinusoid)], dim=1) - return x - - -@torch.amp.autocast(device_type="cuda", enabled=False) def rope_params(max_seq_len, dim, theta=10000, offset=0): assert dim % 2 == 0 freqs = torch.outer( @@ -245,7 +38,6 @@ def rope_params(max_seq_len, dim, theta=10000, offset=0): return freqs -@torch.amp.autocast(device_type="cuda", enabled=False) def rope_apply(x, grid_sizes, freqs, time_stride=1): n, c = x.size(2), x.size(3) // 2 @@ -291,343 +83,445 @@ def pad_freqs(original_tensor, target_len): return padded_tensor -class RMSNorm(nn.Module): - def __init__(self, dim, eps=1e-5): - super().__init__() - self.dim = dim - self.eps = eps - self.weight = nn.Parameter(torch.ones(dim)) +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 - def forward(self, x): - return self._norm(x.float()).type_as(x) * self.weight + 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 _norm(self, x): - return x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps) +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 LayerNorm(nn.LayerNorm): - """ - LayerNorm without learnable affine parameters. + +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, dim, eps=1e-6, elementwise_affine=False): - super().__init__(dim, elementwise_affine=elementwise_affine, eps=eps) + def __init__(self): + self.key: torch.Tensor | None = None + self.value: torch.Tensor | None = None - def forward(self, x): - return super().forward(x.float()).type_as(x) + 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 -class SelfAttention(nn.Module): - def __init__( - self, - dim, - num_heads, - window_size=(-1, -1), - qk_norm=True, - eps=1e-6, - ): - assert dim % num_heads == 0 - super().__init__() - self.dim = dim - self.num_heads = num_heads - self.head_dim = dim // num_heads - self.window_size = window_size - self.qk_norm = qk_norm - self.eps = eps + def clear(self): + self.key = None + self.value = None - # layers - self.q = nn.Linear(dim, dim) - self.k = nn.Linear(dim, dim) - self.v = nn.Linear(dim, dim) - self.o = nn.Linear(dim, dim) - self.norm_q = RMSNorm(dim, eps=eps) if qk_norm else nn.Identity() - self.norm_k = RMSNorm(dim, eps=eps) if qk_norm else nn.Identity() - - def forward(self, *args, method, **kwargs): - return getattr(self, method)(*args, **kwargs) - - def pre_attention(self, x): - b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim - - # query, key, value function - def qkv_fn(x): - q = self.norm_q(self.q(x)).view(b, s, n, d) - k = self.norm_k(self.k(x)).view(b, s, n, d) - v = self.v(x).view(b, s, n, d) - return q, k, v - - q, k, v = qkv_fn(x) - - return q, k, v - - def post_attention(self, x): - # output - x = x.flatten(2) - x = self.o(x) - return x +class WanAnimate2KVCache: + """Container holding one [`WanAnimate2KVLayerCache`] per transformer layer.""" -class CrossAttention(SelfAttention): - def __init__(self, dim, num_heads, window_size=(-1, -1), qk_norm=True, eps=1e-6, use_img_emb=True): - super().__init__(dim, num_heads, window_size, qk_norm, eps) - self.use_img_emb = use_img_emb - if use_img_emb: - self.k_img = nn.Linear(dim, dim) - self.v_img = nn.Linear(dim, dim) - self.norm_k_img = RMSNorm(dim, eps=eps) if qk_norm else nn.Identity() + def __init__(self, num_layers: int): + self.layer_caches = [WanAnimate2KVLayerCache() for _ in range(num_layers)] - def forward(self, x, context, context_lens, counter=0): - """ - x: [B, L1, C]. - context: [B, L2, C]. - context_lens: [B]. - """ - if self.use_img_emb: - context_img = context[:, :257] - context = context[:, 257:] - else: - context = context - - b, n, d = x.size(0), self.num_heads, self.head_dim - - # compute query, key, value - q = self.norm_q(self.q(x)).view(b, -1, n, d) - k = self.norm_k(self.k(context)).view(b, -1, n, d) - v = self.v(context).view(b, -1, n, d) - - if self.use_img_emb: - k_img = self.norm_k_img(self.k_img(context_img)).view(b, -1, n, d) - v_img = self.v_img(context_img).view(b, -1, n, d) - img_x = flash_attention(q, k_img, v_img, k_lens=None) - # compute attention - x = flash_attention(q, k, v, k_lens=context_lens) - - # output - x = x.flatten(2) - if self.use_img_emb: - img_x = img_x.flatten(2) - x = x + img_x - x = self.o(x) - return x + 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 AttentionBlock(nn.Module): - def __init__( + +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, - dim, - ffn_dim, - num_heads, - window_size=(-1, -1), - qk_norm=True, - cross_attn_norm=False, - eps=1e-6, - use_img_emb=True, - ): - super().__init__() - self.dim = dim - self.ffn_dim = ffn_dim - self.num_heads = num_heads - self.window_size = window_size - self.qk_norm = qk_norm - self.cross_attn_norm = cross_attn_norm - self.eps = eps + 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}.") - # layers - self.norm1 = LayerNorm(dim, eps) + 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 - self.self_attn = SelfAttention(dim, num_heads, window_size, qk_norm, eps) - self.norm3 = LayerNorm(dim, eps, elementwise_affine=True) if cross_attn_norm else nn.Identity() +class WanAnimate2CrossAttnProcessor: + r""" + Cross-attention to the text embeddings, plus an additive branch over the CLIP image embeddings. - self.cross_attn = CrossAttention(dim, num_heads, (-1, -1), qk_norm, eps, use_img_emb=use_img_emb) + 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. + """ - self.norm2 = LayerNorm(dim, eps) - self.ffn = nn.Sequential( - nn.Linear(dim, ffn_dim), - nn.GELU(approximate="tanh"), - nn.Linear(ffn_dim, dim), + _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, ) - # modulation - self.modulation = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5) + hidden_states = hidden_states.flatten(2, 3).type_as(query) - def forward(self, *args, method, **kwargs): - return getattr(self, method)(*args, **kwargs) + 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) - def pre_self_attention(self, x, e): - assert e.dtype == torch.float32 - with torch.amp.autocast(device_type="cuda", dtype=torch.float32): - e = (self.modulation + e).chunk(6, dim=1) - assert e[0].dtype == torch.float32 + key_image = key_image.unflatten(2, (attn.heads, -1)) + value_image = value_image.unflatten(2, (attn.heads, -1)) - q, k, v = self.self_attn(self.norm1(x).float() * (1 + e[1]) + e[0], method="pre_attention") - return q, k, v, e + 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) - def post_self_attention(self, x): - x = self.self_attn(x, method="post_attention") - return x + hidden_states = attn.to_out[0](hidden_states) + hidden_states = attn.to_out[1](hidden_states) + return hidden_states - def cross_attention(self, x, context, context_lens, e): - x = x + self.cross_attn(self.norm3(x), context, context_lens) - y = self.ffn(self.norm2(x).float() * (1 + e[4]) + e[3]) - with torch.amp.autocast(device_type="cuda", dtype=torch.float32): - x = x + y * e[5] - return x +class WanAnimate2Attention(torch.nn.Module, AttentionModuleMixin): + _default_processor_cls = WanAnimate2AttnProcessor + _available_processors = [WanAnimate2AttnProcessor, WanAnimate2CrossAttnProcessor] -class IncontextAttentionBlock(nn.Module): + 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, - window_size=(-1, -1), - qk_norm=True, cross_attn_norm=False, eps=1e-6, refer_stride=1, use_img_emb=True, - sparse_type=0, - log_scale=0.0, ): super().__init__() - self.dim = dim - self.ffn_dim = ffn_dim - self.num_heads = num_heads - self.window_size = window_size - self.qk_norm = qk_norm - self.cross_attn_norm = cross_attn_norm - self.eps = eps self.refer_stride = refer_stride - self.sparse_type = sparse_type - self.log_scale = log_scale - self.block = AttentionBlock( - dim, ffn_dim, num_heads, window_size, qk_norm, cross_attn_norm, eps, use_img_emb=use_img_emb + # 1. Self-attention + self.norm1 = FP32LayerNorm(dim, eps, elementwise_affine=False) + self.self_attn = WanAnimate2Attention( + dim=dim, + heads=num_heads, + eps=eps, + processor=WanAnimate2AttnProcessor(), ) - def forward(self, *args, method, **kwargs): - return getattr(self, method)(*args, **kwargs) - - def forward_ref(self, x_ref, index, k_cache, v_cache, context_ref, freqs_ref, grid_sizes_ref, e_ref, context_lens): - q_ref, k_ref, v_ref, e_ref = self.block(x_ref, e_ref, method="pre_self_attention") - - k_cache[index] = k_ref - v_cache[index] = v_ref - q_ref_add_rope = rope_apply(q_ref, grid_sizes_ref, freqs_ref, self.refer_stride) - k_ref_add_rope = rope_apply(k_ref, grid_sizes_ref, freqs_ref, self.refer_stride) - - ref_f, ref_h, ref_w = grid_sizes_ref[0].tolist() - ref_vail_len = ref_f * ref_h * ref_w - - xout_ref = flash_attention( - q=q_ref_add_rope, - k=k_ref_add_rope, - v=v_ref, - k_lens=torch.tensor([ref_vail_len], dtype=torch.long), - window_size=self.window_size, + # 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(), ) - y_ref = self.block(xout_ref, method="post_self_attention") - - with torch.amp.autocast(device_type="cuda", dtype=torch.float32): - x_ref = x_ref + y_ref * e_ref[2] - - x_ref = self.block(x_ref, context_ref, context_lens, e_ref, method="cross_attention") + # 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), + ) - return x_ref + self.modulation = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5) - def forward_gen( + def forward( self, - x, - index, - k_cache, - v_cache, - block_mask, - context, - freqs, - freqs_ref, - grid_sizes, - grid_sizes_ref, - origin_len, - origin_area, - e, - context_lens, - ): - origin_latent_f = origin_len // 4 + 1 - origin_latent_hw = origin_area[0] * origin_area[1] // 256 - origin_max_len = (origin_latent_f + 1) * origin_latent_hw - origin_ref_max_len = origin_latent_f * origin_latent_hw - - f, h, w = grid_sizes[0].tolist() - vail_len = f * h * w - hw = h * w - - ref_f, ref_h, ref_w = grid_sizes_ref[0].tolist() - ref_vail_len = ref_f * ref_h * ref_w - ref_hw = ref_h * ref_w - - q, k, v, e = self.block(x, e, method="pre_self_attention") - - q = rope_apply(q, grid_sizes, freqs) - k = rope_apply(k, grid_sizes, freqs) - k_ref, v_ref = k_cache[index], v_cache[index] - k_ref = rope_apply(k_ref, grid_sizes_ref, freqs_ref, self.refer_stride) - - B, _, N, C = q.shape - device, dtype = q.device, q.dtype - - target_q_len = math.ceil(origin_max_len / 128) * 128 - target_ref_len = math.ceil(origin_ref_max_len / 128) * 128 - target_kv_len = target_q_len + target_ref_len - - q_padding = q[:, vail_len:].clone() - - q_incontext = torch.zeros(B, target_q_len, N, C, device=device, dtype=dtype) - k_incontext = torch.zeros(B, target_kv_len, N, C, device=device, dtype=dtype) - v_incontext = torch.zeros(B, target_kv_len, N, C, device=device, dtype=dtype) - - q_src = q[:, :vail_len].view(B, f, hw, N, C) - k_src = k[:, :vail_len].view(B, f, hw, N, C) - v_src = v[:, :vail_len].view(B, f, hw, N, C) - - q_incontext[:, : f * origin_latent_hw].view(B, f, origin_latent_hw, N, C)[:, :, :hw] = q_src - k_incontext[:, : f * origin_latent_hw].view(B, f, origin_latent_hw, N, C)[:, :, :hw] = k_src - v_incontext[:, : f * origin_latent_hw].view(B, f, origin_latent_hw, N, C)[:, :, :hw] = v_src - - k_ref_src = k_ref[:, :ref_vail_len].view(B, ref_f, ref_hw, N, C) - v_ref_src = v_ref[:, :ref_vail_len].view(B, ref_f, ref_hw, N, C) - - k_incontext[:, target_q_len : target_q_len + ref_f * origin_latent_hw].view(B, ref_f, origin_latent_hw, N, C)[ - :, :, :ref_hw - ] = k_ref_src - v_incontext[:, target_q_len : target_q_len + ref_f * origin_latent_hw].view(B, ref_f, origin_latent_hw, N, C)[ - :, :, :ref_hw - ] = v_ref_src - - score_mod = _get_score_mod(hw=int(origin_latent_hw), log_scale=self.log_scale) - - xout_full = flex_attention( - q=q_incontext, - k=k_incontext, - v=v_incontext, - block_mask=block_mask, - kernel_options=None, - score_mod=score_mod, + 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) - xout_valid = xout_full[:, : f * origin_latent_hw] - xout_valid = xout_valid.view(B, f, origin_latent_hw, N, C) - xout_vail = xout_valid[:, :, :hw] - xout_vail = xout_vail.reshape(B, f * hw, N, C) - xout = torch.cat([xout_vail, q_padding], dim=1) - - y = self.block(xout, method="post_self_attention") + # 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, + ) - with torch.amp.autocast(device_type="cuda", dtype=torch.float32): - x = x + y * e[2] + # 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) - x = self.block(x, context, context_lens, e, method="cross_attention") - return x + return hidden_states class Head(nn.Module): @@ -640,17 +534,15 @@ def __init__(self, dim, out_dim, patch_size, eps=1e-6): # layers out_dim = math.prod(patch_size) * out_dim - self.norm = LayerNorm(dim, eps) + 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): - assert e.dtype == torch.float32 - with torch.amp.autocast(device_type="cuda", dtype=torch.float32): - e = (self.modulation + e.unsqueeze(1)).chunk(2, dim=1) - x = self.head(self.norm(x) * (1 + e[1]) + e[0]) + 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 @@ -671,14 +563,16 @@ def forward(self, image_embeds): return clip_extra_context_tokens -class WanAnimate2Transformer3DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin): +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 KV cache: a reference video is first encoded - (``forward_ref``) to cache K/V tensors, then the generation forward (``forward_gen``) uses the cached - K/V with a block mask (``flex_attention``) and score modification (``log_scale``) for frame-level - sparse in-context attention. + 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)`): @@ -701,10 +595,6 @@ class WanAnimate2Transformer3DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, F The number of attention heads. num_layers (`int`, defaults to `40`): The number of layers of transformer blocks to use. - window_size (`tuple[int]`, defaults to `(-1, -1)`): - Window size for local attention (-1 indicates global attention). - qk_norm (`bool`, defaults to `True`): - Enable query/key normalization. cross_attn_norm (`bool`, defaults to `True`): Enable cross-attention normalization. eps (`float`, defaults to `1e-6`): @@ -719,16 +609,13 @@ class WanAnimate2Transformer3DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, F 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. - sparse_type (`int`, defaults to `0`): - Sparse attention type. - log_scale (`float`, defaults to `0.0`): - Log scale for score modification in in-context attention. """ _supports_gradient_checkpointing = True _skip_layerwise_casting_patterns = ["patch_embedding", "img_emb", "norm"] - _no_split_modules = ["IncontextAttentionBlock"] - _repeated_blocks = ["IncontextAttentionBlock"] + _no_split_modules = ["WanAnimate2TransformerBlock"] + _repeated_blocks = ["WanAnimate2TransformerBlock"] + _skip_keys = ["kv_cache"] _keep_in_fp32_modules = [ "time_embedding", "time_projection", @@ -752,8 +639,6 @@ def __init__( out_dim: int = 16, num_heads: int = 40, num_layers: int = 40, - window_size: tuple = (-1, -1), - qk_norm: bool = True, cross_attn_norm: bool = True, eps: float = 1e-6, use_img_emb: bool = True, @@ -761,8 +646,6 @@ def __init__( refer_offset_h: int = 0, refer_offset_w: int = -1, refer_stride: int = 1, - sparse_type: int = 0, - log_scale: float = 0.0, ): super().__init__() self.patch_size = patch_size @@ -775,8 +658,6 @@ def __init__( self.out_dim = out_dim self.num_heads = num_heads self.num_layers = num_layers - self.window_size = window_size - self.qk_norm = qk_norm self.cross_attn_norm = cross_attn_norm self.eps = eps self.use_img_emb = use_img_emb @@ -784,8 +665,6 @@ def __init__( self.refer_offset_h = refer_offset_h self.refer_offset_w = refer_offset_w self.refer_stride = refer_stride - self.sparse_type = sparse_type - self.log_scale = log_scale # [Denoising Transformer] # embeddings @@ -796,6 +675,7 @@ def __init__( 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(), @@ -809,18 +689,14 @@ def __init__( # blocks self.blocks = nn.ModuleList( [ - IncontextAttentionBlock( + WanAnimate2TransformerBlock( dim, ffn_dim, num_heads, - window_size, - qk_norm, cross_attn_norm, eps, refer_stride, use_img_emb=use_img_emb, - sparse_type=sparse_type, - log_scale=log_scale, ) for _ in range(num_layers) ] @@ -832,16 +708,11 @@ def __init__( if use_img_emb: self.img_emb = MLPProj(1280, dim) - # initialize weights - self.init_weights() self.gradient_checkpointing = False self.block_masks = {} self.block_mask_grid_sizes = {} - def create_mask(self, origin_len, origin_area, device): - origin_latent_f = origin_len // 4 + 1 - hw = int(np.prod(origin_area).item() // 256) - + def create_mask(self, origin_latent_f, hw, device): q_len = (origin_latent_f + 1) * hw k_len = origin_latent_f * hw @@ -885,41 +756,89 @@ def attention_mask_logic(b, h, q_idx, kv_idx): ) return block_mask - def forward(self, *args, method, **kwargs): - return getattr(self, method)(*args, **kwargs) - - def forward_ref( + def forward( self, - x_ref, - grid_sizes, - k_cache, - v_cache, - clip_fea_ref, - y_ref, - context_ref, - seq_len_ref, - t, - ): + 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*), origin_area (`list[int]`, *optional*): + Frame count and spatial size 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. + """ + 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 - # [reference] - x_ref = [torch.cat([u, v], dim=0) for u, v in zip(x_ref, y_ref)] - # embeddings - x_ref = [self.patch_embedding(u.unsqueeze(0)) for u in x_ref] - grid_sizes_ref = torch.stack([torch.tensor(u.shape[2:], dtype=torch.long) for u in x_ref]) - x_ref = [u.flatten(2).transpose(1, 2) for u in x_ref] - seq_lens_ref = torch.tensor([u.size(1) for u in x_ref], dtype=torch.long) - assert seq_lens_ref.max() <= seq_len_ref - x_ref = torch.cat([torch.cat([u, u.new_zeros(1, seq_len_ref - u.size(1), u.size(2))], dim=1) for u in x_ref]) + + # 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 = grid_sizes[0][0].item() + self.refer_offset_t = offset_grid_sizes[0][0].item() if self.refer_offset_h < 0: - self.refer_offset_h = grid_sizes[0][1].item() + self.refer_offset_h = offset_grid_sizes[0][1].item() if self.refer_offset_w < 0: - self.refer_offset_w = grid_sizes[0][2].item() + self.refer_offset_w = offset_grid_sizes[0][2].item() self.freqs_ref = torch.cat( [ @@ -932,158 +851,85 @@ def forward_ref( if self.freqs_ref.device != device: self.freqs_ref = self.freqs_ref.to(device) - # time embeddings ref - with torch.amp.autocast(device_type="cuda", dtype=torch.float32): - e_ref = self.time_embedding(sinusoidal_embedding_1d(self.freq_dim, t * 0 + 1).float()) - e0_ref = self.time_projection(e_ref).unflatten(1, (6, self.dim)) - assert e_ref.dtype == torch.float32 and e0_ref.dtype == torch.float32 - - # [context_ref] - context_ref = self.text_embedding( - torch.stack([torch.cat([u, u.new_zeros(self.text_len - u.size(0), u.size(1))]) for u in context_ref]) + # 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)) - if self.use_img_emb: - context_clip_ref = self.img_emb(clip_fea_ref) - context_ref = torch.concat([context_clip_ref, context_ref], dim=1) - - context_lens = None - # arguments - kwargs = { - "e_ref": e0_ref, - "grid_sizes_ref": grid_sizes_ref, - "freqs_ref": self.freqs_ref, - "context_ref": context_ref, - "context_lens": context_lens, - } - - for idx, block in enumerate(self.blocks): - if torch.is_grad_enabled() and self.gradient_checkpointing: - x_ref = self._gradient_checkpointing_func( - block.forward_ref, - x_ref, - idx, - k_cache, - v_cache, - **kwargs, - ) - else: - x_ref = block(x_ref, idx, k_cache, v_cache, method="forward_ref", **kwargs) - - def forward_gen( - self, - x, - k_cache, - v_cache, - clip_fea, - y, - context, - seq_len, - t, - grid_sizes_ref, - origin_len, - origin_area, - is_uncondtion=False, - ): - # [denoising] - # params - device = self.patch_embedding.weight.device - x = [torch.cat([u, v], dim=0) for u, v in zip(x, y)] - # embeddings - x = [self.patch_embedding(u.unsqueeze(0)) for u in x] - grid_sizes = torch.stack([torch.tensor(u.shape[2:], dtype=torch.long) for u in x]) - x = [u.flatten(2).transpose(1, 2) for u in x] - seq_lens = torch.tensor([u.size(1) for u in x], dtype=torch.long) - assert seq_lens.max() <= seq_len - x = torch.cat([torch.cat([u, u.new_zeros(1, seq_len - u.size(1), u.size(2))], dim=1) for u in x]) - - assert (self.dim % self.num_heads) == 0 and (self.dim // self.num_heads) % 2 == 0 - d = self.dim // self.num_heads - 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, + 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] + ) ) - if self.freqs.device != device: - self.freqs = self.freqs.to(device) - - if self.refer_offset_t < 0: - self.refer_offset_t = grid_sizes[0][0].item() - if self.refer_offset_h < 0: - self.refer_offset_h = grid_sizes[0][1].item() - if self.refer_offset_w < 0: - self.refer_offset_w = 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) - - # time embeddings - with torch.amp.autocast(device_type="cuda", dtype=torch.float32): - e = self.time_embedding(sinusoidal_embedding_1d(self.freq_dim, t).float()) - e0 = self.time_projection(e).unflatten(1, (6, self.dim)) - assert e.dtype == torch.float32 and e0.dtype == torch.float32 + 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 + ) - # [context] - context_lens = None - context = self.text_embedding( - torch.stack([torch.cat([u, u.new_zeros(self.text_len - u.size(0), u.size(1))]) for u in context]) + 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, ) - if self.use_img_emb: - context_clip = self.img_emb(clip_fea) - context = torch.concat([context_clip, context], dim=1) - - block_mask_id = (origin_len, origin_area[0], origin_area[1]) - if block_mask_id not in self.block_masks: - self.block_masks[block_mask_id] = self.create_mask(origin_len, origin_area, x.device) - block_mask = self.block_masks[block_mask_id] - - # arguments - kwargs = { - "e": e0, - "block_mask": block_mask, - "grid_sizes": grid_sizes, - "freqs": self.freqs, - "context": context, - "grid_sizes_ref": grid_sizes_ref, - "freqs_ref": self.freqs_ref, - "context_lens": context_lens, - "origin_area": origin_area, - "origin_len": origin_len, - } - + # 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: - x = self._gradient_checkpointing_func( - block.forward_gen, - x, - idx, - k_cache, - v_cache, - **kwargs, + hidden_states = self._gradient_checkpointing_func( + block, hidden_states, kv_cache=kv_cache.get(idx), **block_kwargs ) else: - x = block(x, idx, k_cache, v_cache, method="forward_gen", **kwargs) + hidden_states = block(hidden_states, kv_cache=kv_cache.get(idx), **block_kwargs) - # head - x = self.head(x, e) + # 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)] - # unpatchify - x = self.unpatchify(x, grid_sizes) - return [u.float() for u in x] + if not return_dict: + return (output,) + return Transformer2DModelOutput(sample=output) def unpatchify(self, x, grid_sizes): c = self.out_dim @@ -1094,27 +940,3 @@ def unpatchify(self, x, grid_sizes): u = u.reshape(c, *[i * j for i, j in zip(v, self.patch_size)]) out.append(u) return out - - def init_weights(self): - # basic init - for m in self.modules(): - if isinstance(m, nn.Linear): - nn.init.xavier_uniform_(m.weight) - if m.bias is not None: - nn.init.zeros_(m.bias) - - # init embeddings - nn.init.xavier_uniform_(self.patch_embedding.weight.flatten(1)) - for m in self.text_embedding.modules(): - if isinstance(m, nn.Linear): - nn.init.normal_(m.weight, std=0.02) - for m in self.time_embedding.modules(): - if isinstance(m, nn.Linear): - nn.init.normal_(m.weight, std=0.02) - - # init output layer - nn.init.zeros_(self.head.head.weight) - - def load_from_official_state_dict(self, state_dict): - """Load weights from the official Wan-Animate-2 checkpoint.""" - self.load_state_dict(state_dict, strict=True) diff --git a/src/diffusers/modular_pipelines/wan/__init__.py b/src/diffusers/modular_pipelines/wan/__init__.py index 0e0f06297311..284b6c9fa436 100644 --- a/src/diffusers/modular_pipelines/wan/__init__.py +++ b/src/diffusers/modular_pipelines/wan/__init__.py @@ -21,13 +21,11 @@ _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"] = ["WanBlocks"] _import_structure["modular_blocks_wan22"] = ["Wan22Blocks"] _import_structure["modular_blocks_wan22_i2v"] = ["Wan22Image2VideoBlocks"] _import_structure["modular_blocks_wan_i2v"] = ["WanImage2VideoAutoBlocks"] _import_structure["modular_pipeline"] = [ - "WanAnimate2ModularPipeline", "Wan22Image2VideoModularPipeline", "Wan22ModularPipeline", "WanImage2VideoModularPipeline", @@ -44,12 +42,10 @@ from .modular_blocks_wan import WanBlocks from .modular_blocks_wan22 import Wan22Blocks from .modular_blocks_wan22_i2v import Wan22Image2VideoBlocks - from .modular_blocks_wan_animate_2 import WanAnimate2Blocks from .modular_blocks_wan_i2v import WanImage2VideoAutoBlocks from .modular_pipeline import ( Wan22Image2VideoModularPipeline, Wan22ModularPipeline, - WanAnimate2ModularPipeline, WanImage2VideoModularPipeline, WanModularPipeline, ) diff --git a/src/diffusers/modular_pipelines/wan/modular_blocks_wan_animate_2.py b/src/diffusers/modular_pipelines/wan/modular_blocks_wan_animate_2.py deleted file mode 100644 index 21f8cd1bab46..000000000000 --- a/src/diffusers/modular_pipelines/wan/modular_blocks_wan_animate_2.py +++ /dev/null @@ -1,462 +0,0 @@ -# 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 -import torch.nn.functional as F - -from ...configuration_utils import FrozenDict -from ...guiders import ClassifierFreeGuidance -from ...models import WanAnimate2Transformer3DModel -from ...schedulers import DPMSolverMultistepScheduler -from ...utils import logging -from ..modular_pipeline import ( - BlockState, - LoopSequentialPipelineBlocks, - ModularPipelineBlocks, - SequentialPipelineBlocks, -) -from ..modular_pipeline_utils import ComponentSpec, ConfigSpec, InputParam, OutputParam -from .decoders import WanVaeDecoderStep -from .encoders import WanTextEncoderStep -from .modular_pipeline import WanModularPipeline - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - - -def _get_sampling_sigmas(sampling_steps, shift): - sigma = np.linspace(1, 0, sampling_steps + 1)[:sampling_steps] - sigma = shift * sigma / (1 + (shift - 1) * sigma) - return sigma - - -class WanAnimate2ImageEncoderStep(ModularPipelineBlocks): - """Encode reference image with CLIP + VAE, and driving video with VAE.""" - - model_name = "wan" - - @property - def expected_components(self): - return [ - ComponentSpec("image_encoder", None), - ComponentSpec("vae", None), - ComponentSpec("image_processor", None), - ] - - @property - def description(self): - return "Encode reference image (CLIP + VAE) and driving video (VAE) for Wan-Animate-2." - - @property - def inputs(self): - return [ - InputParam("image", required=True, type_hint=object, description="Reference character image."), - InputParam("driving_video", required=True, type_hint=list, description="Driving video frames."), - InputParam("height", required=True, type_hint=int), - InputParam("width", required=True, type_hint=int), - InputParam("prompt_ref", required=False, type_hint=str, default="人物动作的参考视频"), - ] - - @property - def outputs(self): - return [ - OutputParam("clip_fea", type_hint=torch.Tensor, description="CLIP features of reference image."), - OutputParam("ref_latents", type_hint=torch.Tensor, description="VAE latents of reference image."), - OutputParam("condition_latents", type_hint=torch.Tensor, description="VAE latents of driving video."), - ] - - @torch.no_grad() - def __call__(self, components, state): - device = state.device - dtype = components.transformer.dtype - - # CLIP encode reference image - image = components.image_processor(images=state.image, return_tensors="pt").to(device) - image_embeds = components.image_encoder(**image, output_hidden_states=True) - state.clip_fea = image_embeds.hidden_states[-2].to(dtype) - - # VAE encode reference image - ref_pixels = components.image_processor(images=state.image, return_tensors="pt").to( - device=device, dtype=components.vae.dtype - ) - ref_latents = components.vae.encode(ref_pixels) - latents_mean = torch.tensor(components.vae.config.latents_mean).view(1, -1, 1, 1, 1).to(ref_latents) - latents_recip_std = 1.0 / torch.tensor(components.vae.config.latents_std).view(1, -1, 1, 1, 1).to(ref_latents) - state.ref_latents = (ref_latents - latents_mean) * latents_recip_std - - # VAE encode driving video - driving_pixels = state.driving_video.to(device=device, dtype=components.vae.dtype) - condition_latents = components.vae.encode(driving_pixels) - state.condition_latents = (condition_latents - latents_mean) * latents_recip_std - - return components, state - - -class WanAnimate2SetTimestepsStep(ModularPipelineBlocks): - """Set timesteps using custom sigma computation for flow matching.""" - - model_name = "wan" - - @property - def expected_components(self): - return [ComponentSpec("scheduler", DPMSolverMultistepScheduler)] - - @property - def description(self): - return "Set timesteps for Wan-Animate-2 with custom sigmas." - - @property - def inputs(self): - return [ - InputParam("num_inference_steps", required=True, type_hint=int), - InputParam("sample_shift", required=False, type_hint=float, default=5.0), - ] - - @property - def outputs(self): - return [OutputParam("timesteps", type_hint=torch.Tensor)] - - @torch.no_grad() - def __call__(self, components, state): - sigmas = _get_sampling_sigmas(state.num_inference_steps, state.sample_shift) - components.scheduler.set_timesteps(sigmas=sigmas, device=state.device) - state.timesteps = components.scheduler.timesteps - return components, state - - -class WanAnimate2PrepareLatentsStep(ModularPipelineBlocks): - """Prepare noise latents and encode reference (forward_ref -> KV cache).""" - - model_name = "wan" - - @property - def expected_components(self): - return [ComponentSpec("transformer", WanAnimate2Transformer3DModel)] - - @property - def description(self): - return "Prepare noise latents and encode reference video to cache KV." - - @property - def inputs(self): - return [ - InputParam("ref_latents", required=True, type_hint=torch.Tensor), - InputParam("clip_fea_ref", required=True, type_hint=torch.Tensor), - InputParam("condition_latents", required=True, type_hint=torch.Tensor), - InputParam("prompt_ref_embeds", required=True, type_hint=torch.Tensor), - InputParam("height", required=True, type_hint=int), - InputParam("width", required=True, type_hint=int), - InputParam("clip_len", required=True, type_hint=int), - InputParam("generator", required=False, type_hint=torch.Generator), - InputParam("timesteps", required=True, type_hint=torch.Tensor), - ] - - @property - def outputs(self): - return [ - OutputParam("latents", type_hint=torch.Tensor), - OutputParam("k_cache", type_hint=dict), - OutputParam("v_cache", type_hint=dict), - OutputParam("grid_sizes_ref", type_hint=torch.Tensor), - ] - - @torch.no_grad() - def __call__(self, components, state): - device = state.device - dtype = components.transformer.dtype - - latent_h = state.height // 8 - latent_w = state.width // 8 - clip_len = state.clip_len - lat_t = (clip_len + 3) // 4 + 1 - - # Prepare noise - noise = torch.randn( - 16, lat_t, latent_h, latent_w, device=device, dtype=torch.float32, generator=state.generator - ) - state.latents = [noise] - - # Prepare grid sizes for reference - ref_shape = state.condition_latents.shape[2:] - state.grid_sizes_ref = torch.tensor([ref_shape], dtype=torch.long) - - # KV cache - state.k_cache = {} - state.v_cache = {} - - # Reference encoding (forward_ref) - max_seq_len_ref = int(math.ceil(np.prod(ref_shape))) - - # Prepare y_ref (mask + ref_latents) - mask_ref = torch.zeros(1, 4, *state.ref_latents.shape[2:], device=device, dtype=dtype) - mask_ref[:, :, 0:1] = 1 - mask_ref = mask_ref.view(1, -1, 4, *state.ref_latents.shape[2:]).transpose(1, 2).squeeze(0) - y_ref = torch.cat([mask_ref, state.ref_latents[0]], dim=0) - - t_ref = torch.tensor([state.timesteps[0].item()], device=device, dtype=dtype) - - components.transformer( - [state.ref_latents[0]], - grid_sizes=state.grid_sizes_ref, - k_cache=state.k_cache, - v_cache=state.v_cache, - clip_fea_ref=state.clip_fea_ref, - y_ref=[y_ref], - context_ref=[state.prompt_ref_embeds[0]], - seq_len_ref=max_seq_len_ref, - t=t_ref, - method="forward_ref", - ) - - return components, state - - -class WanAnimate2LoopBeforeDenoiser(ModularPipelineBlocks): - """Prepare latent model input for the denoiser.""" - - model_name = "wan" - - @property - def description(self): - return "Prepare latent model input within the denoising loop." - - @property - def inputs(self): - return [InputParam("latents", required=True, type_hint=list)] - - @torch.no_grad() - def __call__(self, components, state, i, t): - state.latent_model_input = state.latents[0] - return components, state - - -class WanAnimate2LoopDenoiser(ModularPipelineBlocks): - """Denoiser that calls forward_gen with cached KV.""" - - model_name = "wan" - - def __init__(self, guider_input_fields=None): - if guider_input_fields is None: - guider_input_fields = {"context": ("prompt_embeds", "negative_prompt_embeds")} - self._guider_input_fields = guider_input_fields - super().__init__() - - @property - def expected_components(self): - return [ - ComponentSpec( - "guider", - ClassifierFreeGuidance, - config=FrozenDict({"guidance_scale": 3.0}), - default_creation_method="from_config", - ), - ComponentSpec("transformer", WanAnimate2Transformer3DModel), - ] - - @property - def description(self): - return "Denoiser step that calls forward_gen with cached KV for Wan-Animate-2." - - @property - def inputs(self): - inputs = [InputParam("num_inference_steps", required=True, type_hint=int)] - guider_names = [] - for v in self._guider_input_fields.values(): - if isinstance(v, tuple): - guider_names.extend(v) - else: - guider_names.append(v) - for name in guider_names: - inputs.append(InputParam(name=name, required=True, type_hint=torch.Tensor)) - return inputs - - @torch.no_grad() - def __call__(self, components, state, i, t): - components.guider.set_state(step=i, num_inference_steps=state.num_inference_steps, timestep=t) - guider_state = components.guider.prepare_inputs_from_block_state(state, self._guider_input_fields) - - for batch in guider_state: - components.guider.prepare_models(components.transformer) - cond_kwargs = batch.as_dict() - cond_kwargs = { - k: v.to(state.dtype) if isinstance(v, torch.Tensor) else v - for k, v in cond_kwargs.items() - if k in self._guider_input_fields - } - - is_uncond = batch.guidance_identifier == "pred_uncond" - batch.noise_pred = components.transformer( - state.latents, - k_cache=state.k_cache, - v_cache=state.v_cache, - clip_fea=state.clip_fea, - y=state.y, - seq_len=state.max_seq_len, - t=t.expand(1), - grid_sizes_ref=state.grid_sizes_ref, - origin_len=state.origin_len, - origin_area=state.origin_area, - method="forward_gen", - is_uncondtion=is_uncond, - **cond_kwargs, - ) - if isinstance(batch.noise_pred, list): - batch.noise_pred = batch.noise_pred[0] - components.guider.cleanup_models(components.transformer) - - state.noise_pred = components.guider(guider_state)[0] - return components, state - - -class WanAnimate2DenoiseLoopWrapper(LoopSequentialPipelineBlocks): - """Denoise loop for Wan-Animate-2: before_denoiser -> denoiser -> after_denoiser (scheduler step).""" - - model_name = "wan" - sub_blocks = [WanAnimate2LoopBeforeDenoiser, WanAnimate2LoopDenoiser] - - @property - def description(self): - return "Denoise loop for Wan-Animate-2 using forward_gen with cached KV." - - @property - def inputs(self): - return [ - InputParam("latents", required=True, type_hint=list), - InputParam("k_cache", required=True, type_hint=dict), - InputParam("v_cache", required=True, type_hint=dict), - InputParam("timesteps", required=True, type_hint=torch.Tensor), - ] - - @property - def outputs(self): - return [OutputParam("latents", type_hint=torch.Tensor)] - - @torch.no_grad() - def after_denoiser(self, components, state, i, t): - temp_x0 = components.scheduler.step( - state.noise_pred.unsqueeze(0), - t, - state.latents[0].unsqueeze(0), - return_dict=False, - )[0] - state.latents[0] = temp_x0.squeeze(0) - return components, state - - -# ==================== -# 1. CORE DENOISE -# ==================== - - -# auto_docstring -class WanAnimate2CoreDenoiseStep(SequentialPipelineBlocks): - """ - Core denoise block for Wan-Animate-2: set_timesteps -> prepare_latents (with ref encoding) -> denoise loop. - - Components: - transformer (`WanAnimate2Transformer3DModel`) scheduler (`DPMSolverMultistepScheduler`) guider - (`ClassifierFreeGuidance`) - - Inputs: - num_inference_steps (`int`): Number of denoising steps. - sample_shift (`float`): Shift for sigma computation. - ref_latents (`Tensor`): VAE latents of reference image. - condition_latents (`Tensor`): VAE latents of driving video. - clip_fea (`Tensor`): CLIP features of reference image. - clip_fea_ref (`Tensor`): CLIP features of driving video. - prompt_embeds (`Tensor`): Text embeddings. - negative_prompt_embeds (`Tensor`): Negative text embeddings. - prompt_ref_embeds (`Tensor`): Reference text embeddings. - height (`int`): Output height. - width (`int`): Output width. - clip_len (`int`): Frames per segment. - generator (`Generator`): Random generator. - - Outputs: - latents (`Tensor`): Denoised latents. - """ - - model_name = "wan" - block_classes = [ - WanAnimate2SetTimestepsStep, - WanAnimate2PrepareLatentsStep, - WanAnimate2DenoiseLoopWrapper, - ] - block_names = ["set_timesteps", "prepare_latents", "denoise"] - - @property - def description(self): - return "Core denoise block for Wan-Animate-2." - - @property - def outputs(self): - return [OutputParam.template("latents")] - - -# ==================== -# 2. FULL BLOCKS -# ==================== - - -# auto_docstring -class WanAnimate2Blocks(SequentialPipelineBlocks): - """ - Modular pipeline for character animation using Wan-Animate-2. - - Components: - text_encoder (`UMT5EncoderModel`) tokenizer (`AutoTokenizer`) image_encoder (`CLIPVisionModel`) - transformer (`WanAnimate2Transformer3DModel`) scheduler (`DPMSolverMultistepScheduler`) guider - (`ClassifierFreeGuidance`) vae (`AutoencoderKLWan`) video_processor (`VideoProcessor`) - - Inputs: - prompt (`str`): Text prompt describing the character. - negative_prompt (`str`): Negative prompt. - prompt_ref (`str`): Reference prompt for driving video. - image (`PIL.Image`): Reference character image. - driving_video (`list`): Driving video frames. - height (`int`): Output height. - width (`int`): Output width. - clip_len (`int`): Frames per segment. - num_inference_steps (`int`): Number of denoising steps. - sample_shift (`float`): Shift for sigma computation. - generator (`Generator`): Random generator. - output_type (`str`): Output format. - - Outputs: - videos (`list`): The generated videos. - """ - - model_name = "wan" - block_classes = [ - WanTextEncoderStep, - WanAnimate2ImageEncoderStep, - WanAnimate2CoreDenoiseStep, - WanVaeDecoderStep, - ] - block_names = [ - "text_encoder", - "image_encoder", - "denoise", - "decode", - ] - - @property - def description(self): - return "Modular pipeline for character animation using Wan-Animate-2." - - @property - def outputs(self): - return [OutputParam.template("videos")] diff --git a/src/diffusers/modular_pipelines/wan/modular_pipeline.py b/src/diffusers/modular_pipelines/wan/modular_pipeline.py index 74069843b714..a360440c9251 100644 --- a/src/diffusers/modular_pipelines/wan/modular_pipeline.py +++ b/src/diffusers/modular_pipelines/wan/modular_pipeline.py @@ -139,13 +139,3 @@ class Wan22Image2VideoModularPipeline(Wan22ModularPipeline): """ default_blocks_name = "Wan22Image2VideoBlocks" - - -class WanAnimate2ModularPipeline(WanModularPipeline): - """ - 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" diff --git a/src/diffusers/pipelines/pipeline_utils.py b/src/diffusers/pipelines/pipeline_utils.py index 60e16bdd47d8..a683973df5d7 100644 --- a/src/diffusers/pipelines/pipeline_utils.py +++ b/src/diffusers/pipelines/pipeline_utils.py @@ -30,36 +30,22 @@ import requests import torch from huggingface_hub import ( + DDUFEntry, ModelCard, create_repo, + get_cached_repo_tree, hf_hub_download, model_info, + read_dduf_file, snapshot_download, ) -try: - from huggingface_hub import DDUFEntry, read_dduf_file -except ImportError: - DDUFEntry = None - read_dduf_file = None -try: - from huggingface_hub import get_cached_repo_tree -except ImportError: - get_cached_repo_tree = None -try: - from huggingface_hub.errors import CachedRepoTreeNotFoundError -except ImportError: - class CachedRepoTreeNotFoundError(Exception): - pass +from huggingface_hub.errors import CachedRepoTreeNotFoundError from huggingface_hub.utils import ( HfHubHTTPError, LocalEntryNotFoundError, + OfflineModeIsEnabled, validate_hf_hub_args, ) -try: - from huggingface_hub.utils import OfflineModeIsEnabled -except ImportError: - class OfflineModeIsEnabled(Exception): - pass from packaging import version from tqdm.auto import tqdm from typing_extensions import Self diff --git a/src/diffusers/pipelines/wan/pipeline_wan_animate_2.py b/src/diffusers/pipelines/wan/pipeline_wan_animate_2.py index ef1671441020..86ae8cbbdf4a 100644 --- a/src/diffusers/pipelines/wan/pipeline_wan_animate_2.py +++ b/src/diffusers/pipelines/wan/pipeline_wan_animate_2.py @@ -24,6 +24,7 @@ from ...image_processor import PipelineImageInput from ...loaders import WanLoraLoaderMixin from ...models import AutoencoderKLWan, WanAnimate2Transformer3DModel +from ...models.transformers.transformer_wan_animate_2 import WanAnimate2KVCache from ...schedulers import DPMSolverMultistepScheduler from ...utils import logging from ...video_processor import VideoProcessor @@ -128,8 +129,7 @@ def clip_visual_encode(image_encoder, tensor, device, dtype): 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 - with torch.amp.autocast(device_type="cuda", dtype=dtype): - out = image_encoder(pixel_values=videos, output_hidden_states=True) + out = image_encoder(pixel_values=videos.to(dtype), output_hidden_states=True) return out.hidden_states[-2] @@ -607,25 +607,21 @@ def __call__( "is_uncondtion": True, } - # KV cache - k_cache = {} - v_cache = {} + kv_cache = WanAnimate2KVCache(self.transformer.config.num_layers) # Phase 1: encode reference — cast all inputs to transformer dtype t_ref = torch.tensor([timesteps[0].item()], device=device, dtype=self.transformer.dtype) - with torch.amp.autocast(device_type="cuda", dtype=torch.bfloat16): - self.transformer( - [condition_latents[0].to(self.transformer.dtype)] if condition_latents.ndim == 5 else [condition_latents.to(self.transformer.dtype)], - grid_sizes=grid_sizes_ref, - k_cache=k_cache, - v_cache=v_cache, - clip_fea_ref=arg_ref_c["clip_fea_ref"].to(self.transformer.dtype), - y_ref=[y.to(self.transformer.dtype) for y in arg_ref_c["y_ref"]], - context_ref=[c.to(self.transformer.dtype) for c in arg_ref_c["context_ref"]], - seq_len_ref=max_seq_len_ref, - t=t_ref, - method="forward_ref", - ) + self.transformer( + [condition_latents[0].to(self.transformer.dtype)] if condition_latents.ndim == 5 else [condition_latents.to(self.transformer.dtype)], + timestep=t_ref, + encoder_hidden_states=[c.to(self.transformer.dtype) for c in arg_ref_c["context_ref"]], + encoder_hidden_states_image=arg_ref_c["clip_fea_ref"].to(self.transformer.dtype), + condition_latents=[y.to(self.transformer.dtype) for y in arg_ref_c["y_ref"]], + kv_cache=kv_cache, + kv_cache_mode="extract", + seq_len=max_seq_len_ref, + offset_grid_sizes=grid_sizes_ref, + ) # Phase 2: denoising loop from tqdm import tqdm @@ -633,47 +629,40 @@ def __call__( for i, t in tqdm(enumerate(timesteps), total=len(timesteps), desc=f"Segment {seg_idx+1}/{num_segments}"): timestep = torch.stack([t]) - with torch.amp.autocast(device_type="cuda", dtype=torch.bfloat16): - # Conditional - noise_pred_cond = self.transformer( - latents, - k_cache=k_cache, - v_cache=v_cache, - clip_fea=arg_c["clip_fea"], - y=arg_c["y"], - context=arg_c["context"], + # Conditional + noise_pred_cond = self.transformer( + [l.to(self.transformer.dtype) for l in latents], + timestep=timestep, + encoder_hidden_states=arg_c["context"], + encoder_hidden_states_image=arg_c["clip_fea"], + condition_latents=arg_c["y"], + kv_cache=kv_cache, + kv_cache_mode="cached", + seq_len=max_seq_len, + reference_grid_sizes=grid_sizes_ref, + origin_len=arg_c["origin_len"], + origin_area=arg_c["origin_area"], + ).sample[0] + + if self.do_classifier_free_guidance: + noise_pred_uncond = self.transformer( + [l.to(self.transformer.dtype) for l in latents], + timestep=timestep, + encoder_hidden_states=arg_null["context"], + encoder_hidden_states_image=arg_null["clip_fea"], + condition_latents=arg_null["y"], + kv_cache=kv_cache, + kv_cache_mode="cached", seq_len=max_seq_len, - t=timestep, - grid_sizes_ref=grid_sizes_ref, - origin_len=arg_c["origin_len"], - origin_area=arg_c["origin_area"], - method="forward_gen", - ) - if isinstance(noise_pred_cond, list): - noise_pred_cond = noise_pred_cond[0] - - if self.do_classifier_free_guidance: - noise_pred_uncond = self.transformer( - latents, - k_cache=k_cache, - v_cache=v_cache, - clip_fea=arg_null["clip_fea"], - y=arg_null["y"], - context=arg_null["context"], - seq_len=max_seq_len, - t=timestep, - grid_sizes_ref=grid_sizes_ref, - origin_len=arg_null["origin_len"], - origin_area=arg_null["origin_area"], - method="forward_gen", - is_uncondtion=True, - ) - if isinstance(noise_pred_uncond, list): - noise_pred_uncond = noise_pred_uncond[0] - - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_cond - noise_pred_uncond) - else: - noise_pred = noise_pred_cond + reference_grid_sizes=grid_sizes_ref, + origin_len=arg_null["origin_len"], + origin_area=arg_null["origin_area"], + is_uncondtion=True, + ).sample[0] + + noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_cond - noise_pred_uncond) + else: + noise_pred = noise_pred_cond # Scheduler step temp_x0 = sample_scheduler.step( From 2b722ed623bc4ef0d04778cc36170210b8d6eaf9 Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Sat, 8 Aug 2026 03:31:19 +0000 Subject: [PATCH 03/19] Free per-segment KV cache and latents at the segment boundary Each segment allocates a fresh KV cache holding the reference tokens for every layer -- tens of GB at high resolution. Holding the previous segment's cache alive while the next one is built fragmented the allocator enough to OOM mid-run on an 80GB card. Move finished frames to CPU, clear the cache and drop the per-segment latents before starting the next segment. `out_frames` is deliberately kept: the next segment conditions on its tail. Co-Authored-By: Claude Opus 5 --- src/diffusers/pipelines/wan/pipeline_wan_animate_2.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/src/diffusers/pipelines/wan/pipeline_wan_animate_2.py b/src/diffusers/pipelines/wan/pipeline_wan_animate_2.py index 86ae8cbbdf4a..59aca4be87ed 100644 --- a/src/diffusers/pipelines/wan/pipeline_wan_animate_2.py +++ b/src/diffusers/pipelines/wan/pipeline_wan_animate_2.py @@ -692,16 +692,23 @@ def __call__( if start > 0: out_frames = out_frames[:, :, mask_reft_len:] - all_out_frames.append(out_frames) + all_out_frames.append(out_frames.cpu()) start += effective_segment end += effective_segment + # Each segment allocates a fresh KV cache — at 720p that is tens of GB, and holding + # the previous one while the next is built fragments the allocator enough to OOM. + kv_cache.clear() + # `out_frames` is deliberately kept: the next segment conditions on its tail. + del kv_cache, latents, x0 + torch.cuda.empty_cache() + # Reset scheduler for next segment sample_scheduler.set_timesteps(num_inference_steps, device=device) timesteps = sample_scheduler.timesteps # Concatenate all segments - video = torch.cat(all_out_frames, dim=2)[:, :, :real_frame_len] + video = torch.cat(all_out_frames, dim=2)[:, :, :real_frame_len].to(device) # Remove letterbox padding (crop black borders) p_info = ref_padding_info From 270a84e818bcfe05a869872bcc943e240efb40fa Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Sat, 8 Aug 2026 19:33:19 +0000 Subject: [PATCH 04/19] Let `load_video` report the video's frame rate MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A list of frames does not carry the rate it was sampled at, so a pipeline that has to resample its input to the frame rate the model works at cannot get that number from `load_video` — even though imageio hands it to us and we throw it away. Add an opt-in `return_fps`; GIFs get it from the frame duration. Opt-in, so every existing caller keeps returning a plain list of images. Co-Authored-By: Claude Opus 5 --- src/diffusers/utils/loading_utils.py | 20 +++++++++++++++++--- 1 file changed, 17 insertions(+), 3 deletions(-) diff --git a/src/diffusers/utils/loading_utils.py b/src/diffusers/utils/loading_utils.py index c4fee0cfdd83..7acfd6f6597f 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 From fc0c23bad32effc274101f6768a115f4da6e9fbe Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Sat, 8 Aug 2026 19:33:19 +0000 Subject: [PATCH 05/19] Preprocess Wan-Animate-2 with the image/video processors MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The pipeline reached for decord and cv2 to do what the processors already do. decord was an undeclared dependency imported inside `__call__`, and cv2 is not a diffusers requirement — it is in the deps table but not in `install_requires`, and only consisid and `export_utils` touch it, both behind local imports. `WanAnimateImageProcessor` already letterboxes for `WanAnimatePipeline`, which is what `padding_resize` and `resize_by_area` were hand-rolling: keep the aspect ratio, fill the remainder with black. Add the video counterpart and use both, so the driving video arrives as frames from `load_video` like every other video pipeline takes it, with `driving_video_fps` carrying the one thing a frame list cannot — decord used to read the source rate itself, and the 30 -> 24 resample is load-bearing. Also drop the `seed` argument in favour of the `generator` we already accept, and sample with `self.scheduler` instead of building a scheduler per call from `flow_solver` and `sample_shift`. The registered scheduler was dead weight before; the checkpoints now carry the right one, and swapping it is documented. `_encode_vae` was defined but unused while its body was inlined three times. Preprocessing is equivalent, not identical: output dimensions and resampled frame indices match exactly, and content-aligned the difference is 0.2% of full range — PIL lanczos against cv2 INTER_AREA. PIL also centres the paste one row lower than cv2 did; each path's crop follows its own paste. Seeded end to end on the distilled model that lands at 23.7 dB, the same band the attention refactor already sits in, with an identical contact sheet. Co-Authored-By: Claude Opus 5 --- .../pipelines/wan/image_processor.py | 9 + .../pipelines/wan/pipeline_wan_animate_2.py | 309 ++++++------------ 2 files changed, 109 insertions(+), 209 deletions(-) diff --git a/src/diffusers/pipelines/wan/image_processor.py b/src/diffusers/pipelines/wan/image_processor.py index fa18150fcc6e..261a72d6c42a 100644 --- a/src/diffusers/pipelines/wan/image_processor.py +++ b/src/diffusers/pipelines/wan/image_processor.py @@ -20,6 +20,7 @@ from ...configuration_utils import register_to_config from ...image_processor import VaeImageProcessor from ...utils import PIL_INTERPOLATION +from ...video_processor import VideoProcessor class WanAnimateImageProcessor(VaeImageProcessor): @@ -182,3 +183,11 @@ def get_default_height_width( width = round(np.sqrt(max_area / aspect_ratio)) // mod_value_w * mod_value_w return height, width + + +class WanAnimateVideoProcessor(VideoProcessor, WanAnimateImageProcessor): + r""" + Video counterpart of [`WanAnimateImageProcessor`]. `preprocess_video(..., resize_mode="fill")` letterboxes every + frame into the target frame: the aspect ratio is preserved and the remainder is filled with `fill_color` (black by + default) rather than with stretched image data. + """ diff --git a/src/diffusers/pipelines/wan/pipeline_wan_animate_2.py b/src/diffusers/pipelines/wan/pipeline_wan_animate_2.py index 59aca4be87ed..817fbbfe1823 100644 --- a/src/diffusers/pipelines/wan/pipeline_wan_animate_2.py +++ b/src/diffusers/pipelines/wan/pipeline_wan_animate_2.py @@ -12,52 +12,28 @@ # See the License for the specific language governing permissions and # limitations under the License. -import inspect import math -from typing import Callable +from typing import Any, Callable -import cv2 import numpy as np import torch import torch.nn.functional as F +from tqdm import tqdm from ...image_processor import PipelineImageInput from ...loaders import WanLoraLoaderMixin from ...models import AutoencoderKLWan, WanAnimate2Transformer3DModel from ...models.transformers.transformer_wan_animate_2 import WanAnimate2KVCache -from ...schedulers import DPMSolverMultistepScheduler +from ...schedulers import SchedulerMixin from ...utils import logging -from ...video_processor import VideoProcessor from ..pipeline_utils import DiffusionPipeline +from .image_processor import WanAnimateVideoProcessor from .pipeline_output import WanPipelineOutput logger = logging.get_logger(__name__) # pylint: disable=invalid-name -def get_sampling_sigmas(sampling_steps, shift): - sigma = np.linspace(1, 0, sampling_steps + 1)[:sampling_steps] - sigma = shift * sigma / (1 + (shift - 1) * sigma) - return sigma - - -def retrieve_timesteps(scheduler, num_inference_steps=None, device=None, sigmas=None, **kwargs): - if sigmas is not None: - accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) - if not accept_sigmas: - raise ValueError( - f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" - f" sigmas schedules." - ) - scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs) - timesteps = scheduler.timesteps - num_inference_steps = len(timesteps) - else: - scheduler.set_timesteps(num_inference_steps, device=device, **kwargs) - timesteps = scheduler.timesteps - return timesteps, num_inference_steps - - def get_i2v_mask(lat_t, lat_h, lat_w, mask_len=1, device="cuda"): """Create an i2v mask in latent space. @@ -75,56 +51,19 @@ def get_i2v_mask(lat_t, lat_h, lat_w, mask_len=1, device="cuda"): CLIP_STD = [0.26862954, 0.26130258, 0.27577711] -def get_frame_indices(frame_num, video_fps, clip_length, train_fps): - """Resample video frames to target fps.""" - times = np.arange(0, clip_length) / train_fps +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, frame_num - 1).tolist() - - -def padding_resize(img_ori, height, width, padding_color=(0, 0, 0), interpolation=cv2.INTER_LINEAR): - """Letterbox resize: keep aspect ratio + black padding to exact (height, width).""" - ori_h, ori_w = img_ori.shape[:2] - channel = img_ori.shape[2] if img_ori.ndim == 3 else 1 - img_pad = np.zeros((height, width, channel), dtype=np.uint8) - img_pad[:] = padding_color - - if ori_h / ori_w > height / width: - new_w = int(height / ori_h * ori_w) - img = cv2.resize(img_ori, (new_w, height), interpolation=interpolation) - padding = (width - new_w) // 2 - if img.ndim == 2: - img = img[:, :, np.newaxis] - img_pad[:, padding : padding + new_w, :] = img - return img_pad, {"padding_type": "width", "padding": padding, "side_long": new_w} - else: - new_h = int(width / ori_w * ori_h) - img = cv2.resize(img_ori, (width, new_h), interpolation=interpolation) - padding = (height - new_h) // 2 - if img.ndim == 2: - img = img[:, :, np.newaxis] - img_pad[padding : padding + new_h, :, :] = img - return img_pad, {"padding_type": "height", "padding": padding, "side_long": new_h} - - -def resize_by_area(image, target_area, divisor=16): - """Resize keeping aspect ratio targeting area, pad to exact dims. Returns (image, padding_info).""" - h, w = image.shape[:2] - aspect_ratio = w / h - new_h = math.sqrt(target_area / aspect_ratio) - new_w = target_area / new_h - new_w, new_h = int((new_w // divisor) * divisor), int((new_h // divisor) * divisor) - interpolation = cv2.INTER_AREA if (new_w * new_h < w * h) else cv2.INTER_LINEAR - return padding_resize(image, new_h, new_w, interpolation=interpolation) + return np.clip(frame_indices, 0, num_frames - 1).tolist() 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 = 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) @@ -151,8 +90,10 @@ class WanAnimate2Pipeline(DiffusionPipeline, WanLoraLoaderMixin): CLIP vision model for encoding the reference image. transformer ([`WanAnimate2Transformer3DModel`]): The Wan-Animate-2 transformer model. - scheduler ([`DPMSolverMultistepScheduler`]): - A scheduler for flow matching. + scheduler ([`SchedulerMixin`]): + A flow-matching scheduler to be used in combination with `transformer` to denoise the encoded latents. + The reference implementation samples with `DPMSolverMultistepScheduler` (`flow_shift=5.0`) for the base + model and `FlowMatchEulerDiscreteScheduler` (`shift=5.0`) for the distilled one. vae ([`AutoencoderKLWan`]): The Wan VAE model. """ @@ -165,7 +106,7 @@ def __init__( tokenizer, text_encoder, vae: AutoencoderKLWan, - scheduler: DPMSolverMultistepScheduler, + scheduler: SchedulerMixin, image_encoder, transformer: WanAnimate2Transformer3DModel, ): @@ -182,7 +123,11 @@ def __init__( self.vae_scale_factor_temporal = self.vae.config.scale_factor_temporal if getattr(self, "vae", None) else 4 self.vae_scale_factor_spatial = self.vae.config.scale_factor_spatial if getattr(self, "vae", None) else 8 - self.video_processor = VideoProcessor(vae_scale_factor=self.vae_scale_factor_spatial) + # Wan-Animate-2 letterboxes the reference image and the driving video into the same frame: aspect + # ratio preserved, the remainder filled with black. That is `resize_mode="fill"` with `fill_color=0`. + self.video_processor = WanAnimateVideoProcessor( + vae_scale_factor=self.vae_scale_factor_spatial, spatial_patch_size=(2, 2) + ) def _get_t5_prompt_embeds(self, prompt, device=None, dtype=None, max_sequence_length=512): device = device or self._execution_device @@ -220,10 +165,9 @@ def encode_image(self, image, device=None): image_embeds = self.image_encoder(**processed, output_hidden_states=True) return image_embeds.hidden_states[-2] - def _encode_vae(self, video, device, dtype): - """Encode video to latents using VAE, with standardization.""" - video = video.to(device=device, dtype=dtype) - latents = self.vae.encode(video) + def _encode_vae(self, video): + """VAE-encode a `[B, C, T, H, W]` clip and standardize the latents.""" + latents = self.vae.encode(video.to(self.vae.dtype)) if hasattr(latents, "latent_dist"): latents = latents.latent_dist.mode() elif hasattr(latents, "latents"): @@ -280,7 +224,7 @@ def num_timesteps(self): def __call__( self, image: PipelineImageInput, - driving_video: list, + driving_video: list[Any], prompt: str | list[str] = None, negative_prompt: str | list[str] = None, prompt_ref: str = "人物动作的参考视频", @@ -289,11 +233,9 @@ def __call__( clip_len: int = 81, first_num: int = 1, fps: int = 24, + driving_video_fps: float | None = None, num_inference_steps: int = 40, guidance_scale: float = 3.0, - sample_shift: float = 5.0, - flow_solver: str = "dpm", - seed: int = -1, generator: torch.Generator | list[torch.Generator] | None = None, output_type: str | None = "np", return_dict: bool = True, @@ -307,8 +249,10 @@ def __call__( Args: image (`PipelineImageInput`): The reference character image. - driving_video (`list`): - The driving video (list of PIL images or tensors) that provides motion. + driving_video (`list[PIL.Image.Image]`, `np.ndarray` or `torch.Tensor`): + The driving video that provides the motion, in any format accepted by + [`~video_processor.VideoProcessor.preprocess_video`]. Load one from disk with + [`~utils.load_video`]. prompt (`str` or `list[str]`): The text prompt describing the character appearance and background. negative_prompt (`str` or `list[str]`, *optional*): @@ -316,23 +260,28 @@ def __call__( prompt_ref (`str`, defaults to `"人物动作的参考视频"`): The reference prompt for the driving video context. height (`int`, defaults to `800`): - The height of the generated video. + Together with `width`, the target *area* (`height * width`) of the generated video. The aspect ratio + is taken from `image`, so the video is rarely exactly `height` x `width` — both dimensions are + rescaled to hit that area and then floored to a multiple of 16. width (`int`, defaults to `640`): - The width of the generated video. + See `height`. clip_len (`int`, defaults to `81`): The number of frames in each inference segment. first_num (`int`, defaults to `1`): The number of conditioning frames from the previous segment. fps (`int`, defaults to `24`): - The output video FPS. + The frame rate the model generates at. `driving_video` is resampled to it when + `driving_video_fps` is given. + driving_video_fps (`float`, *optional*): + The frame rate `driving_video` was captured at — a list of frames does not carry it, so + [`~utils.load_video`] will report it with `return_fps=True`. When set, the driving frames are + nearest-neighbour resampled from it to `fps`; when `None` they are used as-is. num_inference_steps (`int`, defaults to `40`): The number of denoising steps. guidance_scale (`float`, defaults to `3.0`): Guidance scale for classifier-free guidance. - sample_shift (`float`, defaults to `5.0`): - The shift parameter for sigma computation. - seed (`int`, defaults to `-1`): - Random seed. -1 means random. + generator (`torch.Generator`, *optional*): + A generator to make generation deterministic. output_type (`str`, defaults to `"np"`): The output format. return_dict (`bool`, defaults to `True`): @@ -344,40 +293,37 @@ def __call__( self._guidance_scale = guidance_scale device = self._execution_device - if seed >= 0: - torch.manual_seed(seed) - torch.cuda.manual_seed_all(seed) - np.random.seed(seed) - - if generator is None: - generator = torch.Generator(device=device) - if seed >= 0: - generator.manual_seed(seed) - - # 2. Preprocess reference image (letterbox resize — do this first to get actual dims) - ref_np = np.array(image) # PIL → numpy [H, W, C] - ref_pad, ref_padding_info = resize_by_area(ref_np, width * height, divisor=16) - actual_h, actual_w = ref_pad.shape[:2] - - # 3. Prepare driving video frames (FPS resampling + letterbox to match ref dims) - import decord - - vr = decord.VideoReader(driving_video if isinstance(driving_video, str) else None) - video_fps = vr.get_avg_fps() - frame_num = len(vr) - target_num = int(frame_num / video_fps * fps) - idxs = get_frame_indices(frame_num, video_fps, target_num, fps) - frames_np = vr.get_batch(idxs).asnumpy() # [T, H, W, C] uint8 - - cond_images_np = [] - for frame in frames_np: - img_pad, _ = padding_resize(frame, actual_h, actual_w) - cond_images_np.append(img_pad) - - driving_video = torch.tensor(np.stack(cond_images_np), dtype=torch.float32) # [T, H, W, C] - driving_video = driving_video / 127.5 - 1.0 # [-1, 1] - driving_video = driving_video.permute(3, 0, 1, 2).unsqueeze(0) # [1, C, T, H, W] - driving_video = driving_video.to(device, dtype=torch.float32) + # 2. Resolve the output frame. `height * width` is a target *area*; the aspect ratio comes from the + # reference image, and both sides are floored to a multiple of `vae_scale_factor_spatial * patch_size` + # so the latent grid divides evenly. + image_height, image_width = self.video_processor.get_default_height_width(image) + mod_value = self.vae_scale_factor_spatial * 2 + aspect_ratio = image_height / image_width + actual_h = int(math.sqrt(height * width * aspect_ratio)) // mod_value * mod_value + actual_w = int(math.sqrt(height * width / aspect_ratio)) // mod_value * mod_value + + # The reference image is letterboxed into that frame. `resize_mode="fill"` keeps the aspect ratio and + # pads the remainder with black; record the pasted box so the bars can be cropped back off the output. + src_w = ( + actual_w if actual_w / actual_h < image_width / image_height else image_width * actual_h // image_height + ) + src_h = ( + actual_h if actual_w / actual_h >= image_width / image_height else image_height * actual_w // image_width + ) + crop_top, crop_left = actual_h // 2 - src_h // 2, actual_w // 2 - src_w // 2 + + image_pixels = self.video_processor.preprocess(image, height=actual_h, width=actual_w, resize_mode="fill").to( + device, dtype=torch.float32 + ) + + # 3. Preprocess the driving video into the same frame, resampling to `fps` first if asked to. + if driving_video_fps is not None: + frame_indices = get_frame_indices(len(driving_video), driving_video_fps, fps) + driving_video = [driving_video[i] for i in frame_indices] + + driving_video = self.video_processor.preprocess_video( + driving_video, height=actual_h, width=actual_w, resize_mode="fill" + ).to(device, dtype=torch.float32) # Pad driving video to be a multiple of (clip_len - first_num) real_frame_len = driving_video.shape[2] @@ -391,6 +337,8 @@ def __call__( # Pad driving video using zigzag (reflect) strategy if num_padding > 0: + # Mirrored real frames, not filler: the model attends to them like any other frame and needs no mask. + # The surplus generated frames are cropped off again with `[:, :, :real_frame_len]` at the end. padding_frames = driving_video[:, :, real_frame_len - num_padding : real_frame_len].flip(2) driving_video = torch.cat([driving_video, padding_frames], dim=2) @@ -409,35 +357,11 @@ def __call__( ) # 5. Encode reference image (VAE + CLIP) - ref_tensor = torch.tensor(ref_pad, dtype=torch.float32) / 127.5 - 1.0 # [-1, 1] - image_pixels = ref_tensor.permute(2, 0, 1).unsqueeze(0).unsqueeze(2).to(device, dtype=torch.float32) - # CLIP features from reference image (direct bicubic to 224×224 from tensor) - clip_fea = clip_visual_encode(self.image_encoder, ref_tensor.permute(2, 0, 1).to(device), device, self.transformer.dtype) - - # VAE encode reference image - ref_pixels = image_pixels.to(self.vae.dtype) - if ref_pixels.ndim == 4: - ref_pixels = ref_pixels.unsqueeze(2) # [B, C, H, W] -> [B, C, 1, H, W] - ref_latents = self.vae.encode(ref_pixels) - if hasattr(ref_latents, "latent_dist"): - ref_latents = ref_latents.latent_dist.mode() - elif hasattr(ref_latents, "latents"): - ref_latents = ref_latents.latents - elif isinstance(ref_latents, (list, tuple)): - ref_latents = torch.stack(ref_latents) if not isinstance(ref_latents[0], torch.Tensor) else ref_latents[0] - latents_mean = ( - torch.tensor(self.vae.config.latents_mean) - .view(1, self.vae.config.z_dim, 1, 1, 1) - .to(ref_latents.device, ref_latents.dtype) - ) - latents_recip_std = 1.0 / torch.tensor(self.vae.config.latents_std).view(1, self.vae.config.z_dim, 1, 1, 1).to( - ref_latents.device, ref_latents.dtype - ) - ref_latents = (ref_latents - latents_mean) * latents_recip_std + clip_fea = clip_visual_encode(self.image_encoder, image_pixels[0], device, self.transformer.dtype) + + ref_latents = self._encode_vae(image_pixels.unsqueeze(2)) # [B, C, H, W] -> [B, C, 1, H, W] - # Derive latent dims from ACTUAL image size after resize_by_area (not requested height/width) - actual_h, actual_w = ref_pad.shape[:2] latent_h = actual_h // self.vae_scale_factor_spatial latent_w = actual_w // self.vae_scale_factor_spatial @@ -449,29 +373,12 @@ def __call__( # CLIP context for reference clip_context = clip_fea - # 5. Set up scheduler - if flow_solver == "euler": - from diffusers import FlowMatchEulerDiscreteScheduler - - sample_scheduler = FlowMatchEulerDiscreteScheduler( - num_train_timesteps=1000, - shift=sample_shift, - use_dynamic_shifting=False, - ) - else: - sample_scheduler = DPMSolverMultistepScheduler.from_config( - self.scheduler.config, - num_train_timesteps=1000, - flow_shift=sample_shift, - use_dynamic_shifting=False, - prediction_type="flow_prediction", - ) - sample_scheduler.set_timesteps(num_inference_steps, device=device) - timesteps = sample_scheduler.timesteps - + # 6. Prepare timesteps + self.scheduler.set_timesteps(num_inference_steps, device=device) + timesteps = self.scheduler.timesteps self._num_timesteps = len(timesteps) - # 6. Segment-based generation loop + # 7. Segment-based generation loop start = 0 end = clip_len all_out_frames = [] @@ -490,16 +397,11 @@ def __call__( else: clip_len_actual = clip_len - # VAE encode the driving video segment - cond_pixels = driving_video[:, :, start : start + clip_len_actual].to(self.vae.dtype) - condition_latents = self.vae.encode(cond_pixels) - if hasattr(condition_latents, "latent_dist"): - condition_latents = condition_latents.latent_dist.mode() - elif hasattr(condition_latents, "latents"): - condition_latents = condition_latents.latents - elif isinstance(condition_latents, (list, tuple)): - condition_latents = condition_latents[0] if isinstance(condition_latents[0], torch.Tensor) else torch.stack(condition_latents) - condition_latents = (condition_latents - latents_mean) * latents_recip_std + # VAE-encode this segment's slice of the driving video. The Wan VAE is causal in time, so + # encoding the whole video once up front and slicing the latents is not the same tensor — + # segments overlap by `first_num` frames and each slice restarts the temporal convolution. + # Encoding per segment is also what a streaming mode would have to do anyway. + condition_latents = self._encode_vae(driving_video[:, :, start : start + clip_len_actual]) # CLIP features from driving video first frame (direct bicubic to 224×224 from tensor) condition_img = driving_video[0, :, 0] # [C, H, W] in [-1, 1] @@ -519,18 +421,11 @@ def __call__( cond_y_input = torch.cat( [prev_frames_interp, torch.zeros(3, T - mask_reft_len - 1, actual_h, actual_w, device=device)], dim=1, - ).to(self.vae.dtype) + ) else: - cond_y_input = torch.zeros(3, T - 1, actual_h, actual_w, device=device).to(self.vae.dtype) - - y_reft = self.vae.encode(cond_y_input.unsqueeze(0)) - if hasattr(y_reft, "latent_dist"): - y_reft = y_reft.latent_dist.mode() - elif hasattr(y_reft, "latents"): - y_reft = y_reft.latents - elif isinstance(y_reft, (list, tuple)): - y_reft = y_reft[0] - y_reft = (y_reft - latents_mean) * latents_recip_std + cond_y_input = torch.zeros(3, T - 1, actual_h, actual_w, device=device) + + y_reft = self._encode_vae(cond_y_input.unsqueeze(0)) if y_reft.ndim == 5: y_reft = y_reft.squeeze(0) # [1, 16, T, H, W] -> [16, T, H, W] @@ -612,7 +507,9 @@ def __call__( # Phase 1: encode reference — cast all inputs to transformer dtype t_ref = torch.tensor([timesteps[0].item()], device=device, dtype=self.transformer.dtype) self.transformer( - [condition_latents[0].to(self.transformer.dtype)] if condition_latents.ndim == 5 else [condition_latents.to(self.transformer.dtype)], + [condition_latents[0].to(self.transformer.dtype)] + if condition_latents.ndim == 5 + else [condition_latents.to(self.transformer.dtype)], timestep=t_ref, encoder_hidden_states=[c.to(self.transformer.dtype) for c in arg_ref_c["context_ref"]], encoder_hidden_states_image=arg_ref_c["clip_fea_ref"].to(self.transformer.dtype), @@ -624,9 +521,7 @@ def __call__( ) # Phase 2: denoising loop - from tqdm import tqdm - - for i, t in tqdm(enumerate(timesteps), total=len(timesteps), desc=f"Segment {seg_idx+1}/{num_segments}"): + for i, t in tqdm(enumerate(timesteps), total=len(timesteps), desc=f"Segment {seg_idx + 1}/{num_segments}"): timestep = torch.stack([t]) # Conditional @@ -665,7 +560,7 @@ def __call__( noise_pred = noise_pred_cond # Scheduler step - temp_x0 = sample_scheduler.step( + temp_x0 = self.scheduler.step( noise_pred.unsqueeze(0), t, latents[0].unsqueeze(0), @@ -703,19 +598,15 @@ def __call__( del kv_cache, latents, x0 torch.cuda.empty_cache() - # Reset scheduler for next segment - sample_scheduler.set_timesteps(num_inference_steps, device=device) - timesteps = sample_scheduler.timesteps + # Each segment is an independent trajectory, so the solver state has to be reset. + self.scheduler.set_timesteps(num_inference_steps, device=device) + timesteps = self.scheduler.timesteps # Concatenate all segments video = torch.cat(all_out_frames, dim=2)[:, :, :real_frame_len].to(device) - # Remove letterbox padding (crop black borders) - p_info = ref_padding_info - if p_info["padding_type"] == "width": - video = video[:, :, :, :, p_info["padding"] : p_info["padding"] + p_info["side_long"]] - else: - video = video[:, :, :, p_info["padding"] : p_info["padding"] + p_info["side_long"], :] + # Crop the reference image's letterbox bars back off + video = video[:, :, :, crop_top : crop_top + src_h, crop_left : crop_left + src_w] video = self.video_processor.postprocess_video(video, output_type=output_type) From 54dabb7663871928f61f9ae4a33bb370dd5af574 Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Mon, 10 Aug 2026 08:14:37 +0000 Subject: [PATCH 06/19] Match Wan-Animate-2 preprocessing to the reference implementation Seed-matched against the approved commit, everything the pipeline hands the transformer is bit-identical -- noise, timesteps, text embeddings, geometry -- except the resized pixels, and the entire pixel residual traced to three resize facts, each verified in isolation (decode and normalization are proven byte-identical): - kernel choice: bilinear matches the driving frames' `INTER_LINEAR` upscale (same filter; only exact-half ties round differently, at most one 8-bit level per pixel), bicubic measures closest to `INTER_AREA` for the reference image's downscale of the PIL kernels diffusers exposes - paste placement: cv2 letterboxes at `(height - src_h) // 2`, PIL at `height // 2 - src_h // 2` -- one row apart when the frame is even and the content odd, previously the largest input difference - the interim `WanAnimateVideoProcessor` inherited `WanAnimateImageProcessor`'s `__init__`, whose bare `super().__init__()` re-registers every shared config field with the parent's defaults -- `resample` was silently lanczos `WanAnimate2VideoProcessor` replaces it: one class, its own `register_to_config` init that does not chain into the decorated parents, the reference paste convention, and per-instance kernels (bicubic for the reference image, bilinear for the driving video). `WanAnimateImageProcessor` and the merged Wan-Animate pipeline are untouched. First-segment output agreement with the approved commit at matched seed is 27.5 dB (base) / 28.9 dB (distilled), above the 25.1 dB the approved commit scores against itself when only the attention backend changes. Divergence in later segments is the chained conditioning amplifying any perturbation, numerical noise included, and is documented where the processors are built. Co-Authored-By: Claude Opus 5 --- .../pipelines/wan/image_processor.py | 55 +++++++++++++++++-- .../pipelines/wan/pipeline_wan_animate_2.py | 26 ++++++--- 2 files changed, 69 insertions(+), 12 deletions(-) diff --git a/src/diffusers/pipelines/wan/image_processor.py b/src/diffusers/pipelines/wan/image_processor.py index 261a72d6c42a..f1a046a515c6 100644 --- a/src/diffusers/pipelines/wan/image_processor.py +++ b/src/diffusers/pipelines/wan/image_processor.py @@ -185,9 +185,56 @@ def get_default_height_width( return height, width -class WanAnimateVideoProcessor(VideoProcessor, WanAnimateImageProcessor): +class WanAnimate2VideoProcessor(VideoProcessor, WanAnimateImageProcessor): r""" - Video counterpart of [`WanAnimateImageProcessor`]. `preprocess_video(..., resize_mode="fill")` letterboxes every - frame into the target frame: the aspect ratio is preserved and the remainder is filled with `fill_color` (black by - default) rather than with stretched image data. + 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). Same letterbox as + [`WanAnimateImageProcessor`], except the resized content is pasted at + `((height - src_h) // 2, (width - src_w) // 2)` -- the placement convention of the reference implementation -- + instead of `(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__`s: they are themselves + # `register_to_config`-decorated, so calling one 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 parents' bodies only validate. + 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 diff --git a/src/diffusers/pipelines/wan/pipeline_wan_animate_2.py b/src/diffusers/pipelines/wan/pipeline_wan_animate_2.py index 817fbbfe1823..834afb2afc38 100644 --- a/src/diffusers/pipelines/wan/pipeline_wan_animate_2.py +++ b/src/diffusers/pipelines/wan/pipeline_wan_animate_2.py @@ -27,7 +27,7 @@ from ...schedulers import SchedulerMixin from ...utils import logging from ..pipeline_utils import DiffusionPipeline -from .image_processor import WanAnimateVideoProcessor +from .image_processor import WanAnimate2VideoProcessor from .pipeline_output import WanPipelineOutput @@ -124,9 +124,19 @@ def __init__( self.vae_scale_factor_temporal = self.vae.config.scale_factor_temporal if getattr(self, "vae", None) else 4 self.vae_scale_factor_spatial = self.vae.config.scale_factor_spatial if getattr(self, "vae", None) else 8 # Wan-Animate-2 letterboxes the reference image and the driving video into the same frame: aspect - # ratio preserved, the remainder filled with black. That is `resize_mode="fill"` with `fill_color=0`. - self.video_processor = WanAnimateVideoProcessor( - vae_scale_factor=self.vae_scale_factor_spatial, spatial_patch_size=(2, 2) + # ratio preserved, the remainder filled with black (`resize_mode="fill"` with `fill_color=0`). + # The reference implementation resizes with cv2, which is not a diffusers dependency, so these + # processors use the closest PIL kernels: bilinear for the driving frames (the same filter as + # `INTER_LINEAR`, but PIL quantizes interpolation weights to 22 bits where cv2 uses 11, so + # exact-half values round in opposite directions -- at most one 8-bit level per pixel) and + # bicubic for the reference image's downscale (`INTER_AREA` has no PIL equivalent; bicubic + # measures closest). Outputs therefore differ very slightly, numerically and visually, from + # the original repository. + self.image_processor_for_reference = WanAnimate2VideoProcessor( + vae_scale_factor=self.vae_scale_factor_spatial, spatial_patch_size=(2, 2), resample="bicubic" + ) + self.video_processor = WanAnimate2VideoProcessor( + vae_scale_factor=self.vae_scale_factor_spatial, spatial_patch_size=(2, 2), resample="bilinear" ) def _get_t5_prompt_embeds(self, prompt, device=None, dtype=None, max_sequence_length=512): @@ -310,11 +320,11 @@ def __call__( src_h = ( actual_h if actual_w / actual_h >= image_width / image_height else image_height * actual_w // image_width ) - crop_top, crop_left = actual_h // 2 - src_h // 2, actual_w // 2 - src_w // 2 + crop_top, crop_left = (actual_h - src_h) // 2, (actual_w - src_w) // 2 - image_pixels = self.video_processor.preprocess(image, height=actual_h, width=actual_w, resize_mode="fill").to( - device, dtype=torch.float32 - ) + image_pixels = self.image_processor_for_reference.preprocess( + image, height=actual_h, width=actual_w, resize_mode="fill" + ).to(device, dtype=torch.float32) # 3. Preprocess the driving video into the same frame, resampling to `fps` first if asked to. if driving_video_fps is not None: From 918c2399a34952ce4dbd754be18ca891b515029f Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Tue, 11 Aug 2026 02:58:37 +0000 Subject: [PATCH 07/19] Add Wan-Animate-2 modular pipeline Ground-up modular decomposition of WanAnimate2Pipeline in modular_pipelines/wan_animate_2/, verified bit-identical to the standard pipeline (tiny fp32 with and without CFG, and the real distilled checkpoint in bf16). - Outer segment loop as LoopSequentialPipelineBlocks (helios style), with the per-segment driving VAE encode, previous-frame conditioning, KV-cache reference extraction, scheduler reset, hand-written denoise loop, and in-loop decode (each segment conditions on the previous segment's decoded pixels) as separate loop blocks. - CFG through the guider; `is_uncondtion` rides the guider's per-branch tuple inputs. Two presets: WanAnimate2Blocks (guidance_scale=3.0) and WanAnimate2DistilledBlocks (guidance_scale=1.0). - Segment-invariant work hoisted out of the loop: text/CLIP encoders (the driving-frame CLIP context is computed once, not per segment), reference VAE encode, and segment geometry. - WanAnimate2VideoProcessor moves into the modular folder; the standard pipeline imports it from there, and pipelines/wan/image_processor.py is back to zero net change. Co-Authored-By: Claude Fable 5 --- src/diffusers/__init__.py | 8 + src/diffusers/modular_pipelines/__init__.py | 12 + .../modular_pipelines/modular_pipeline.py | 1 + .../wan_animate_2/__init__.py | 57 ++ .../wan_animate_2/before_denoise.py | 85 +++ .../wan_animate_2/decoders.py | 121 ++++ .../wan_animate_2/denoise.py | 650 ++++++++++++++++++ .../wan_animate_2/encoders.py | 571 +++++++++++++++ .../modular_blocks_wan_animate_2.py | 125 ++++ .../modular_blocks_wan_animate_2_distilled.py | 125 ++++ .../wan_animate_2/modular_pipeline.py | 71 ++ .../wan_animate_2/video_processor.py | 128 ++++ .../pipelines/wan/image_processor.py | 56 -- .../pipelines/wan/pipeline_wan_animate_2.py | 2 +- .../dummy_torch_and_transformers_objects.py | 60 ++ 15 files changed, 2015 insertions(+), 57 deletions(-) create mode 100644 src/diffusers/modular_pipelines/wan_animate_2/__init__.py create mode 100644 src/diffusers/modular_pipelines/wan_animate_2/before_denoise.py create mode 100644 src/diffusers/modular_pipelines/wan_animate_2/decoders.py create mode 100644 src/diffusers/modular_pipelines/wan_animate_2/denoise.py create mode 100644 src/diffusers/modular_pipelines/wan_animate_2/encoders.py create mode 100644 src/diffusers/modular_pipelines/wan_animate_2/modular_blocks_wan_animate_2.py create mode 100644 src/diffusers/modular_pipelines/wan_animate_2/modular_blocks_wan_animate_2_distilled.py create mode 100644 src/diffusers/modular_pipelines/wan_animate_2/modular_pipeline.py create mode 100644 src/diffusers/modular_pipelines/wan_animate_2/video_processor.py diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py index 34b57dfce51c..36047ec5ca62 100644 --- a/src/diffusers/__init__.py +++ b/src/diffusers/__init__.py @@ -557,6 +557,10 @@ "WanBlocks", "WanImage2VideoAutoBlocks", "WanImage2VideoModularPipeline", + "WanAnimate2Blocks", + "WanAnimate2DistilledBlocks", + "WanAnimate2DistilledModularPipeline", + "WanAnimate2ModularPipeline", "WanModularPipeline", "ZImageAutoBlocks", "ZImageModularPipeline", @@ -1382,6 +1386,10 @@ Wan22Image2VideoBlocks, Wan22Image2VideoModularPipeline, Wan22ModularPipeline, + WanAnimate2Blocks, + WanAnimate2DistilledBlocks, + WanAnimate2DistilledModularPipeline, + WanAnimate2ModularPipeline, WanBlocks, WanImage2VideoAutoBlocks, WanImage2VideoModularPipeline, diff --git a/src/diffusers/modular_pipelines/__init__.py b/src/diffusers/modular_pipelines/__init__.py index 008a654c3fa3..8ec379c2e943 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", @@ -216,6 +222,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 9bf1ddca3b98..7b351be7e527 100644 --- a/src/diffusers/modular_pipelines/modular_pipeline.py +++ b/src/diffusers/modular_pipelines/modular_pipeline.py @@ -132,6 +132,7 @@ 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-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..25dd43ddc0d2 --- /dev/null +++ b/src/diffusers/modular_pipelines/wan_animate_2/before_denoise.py @@ -0,0 +1,85 @@ +# 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 `clip_len` 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("clip_len", type_hint=int, default=81), + InputParam("latent_height", type_hint=int, required=True), + InputParam("latent_width", type_hint=int, required=True), + ] + + @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_segment_frames = (block_state.clip_len - 1) // components.vae_scale_factor_temporal + 1 + ref_shape = [latent_segment_frames, block_state.latent_height, block_state.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, block_state.latent_height // 2, block_state.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..7b52b0f2e886 --- /dev/null +++ b/src/diffusers/modular_pipelines/wan_animate_2/decoders.py @@ -0,0 +1,121 @@ +# 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_top", + required=True, + type_hint=int, + description="Top edge of the reference image content inside the letterboxed frame", + ), + InputParam( + "crop_left", + required=True, + type_hint=int, + description="Left edge of the reference image content inside the letterboxed frame", + ), + InputParam( + "crop_height", + required=True, + type_hint=int, + description="Height of the reference image content inside the letterboxed frame", + ), + InputParam( + "crop_width", + required=True, + type_hint=int, + description="Width of the reference image content inside the letterboxed frame", + ), + 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] + video = video[ + :, + :, + :, + block_state.crop_top : block_state.crop_top + block_state.crop_height, + block_state.crop_left : block_state.crop_left + block_state.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..eac81a5cf7a6 --- /dev/null +++ b/src/diffusers/modular_pipelines/wan_animate_2/denoise.py @@ -0,0 +1,650 @@ +# 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 ..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. The Wan VAE " + "is causal in time, so encoding the whole video once up front and slicing the latents would not be " + "equivalent — each segment restarts the temporal convolution. 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", required=True, type_hint=torch.Tensor), + InputParam("effective_segment", required=True, type_hint=int), + InputParam("clip_len", type_hint=int, default=81), + InputParam("latent_height", required=True, type_hint=int), + InputParam("latent_width", required=True, type_hint=int), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + "condition_latents", + type_hint=torch.Tensor, + description="VAE latents of this segment's driving-video slice", + ), + OutputParam( + "condition_y", + 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 + + start = k * block_state.effective_segment + block_state.condition_latents = encode_vae( + components.vae, block_state.driving_video[:, :, start : start + block_state.clip_len] + ) + + condition_mask = get_i2v_mask( + block_state.condition_latents.shape[2], + block_state.latent_height, + block_state.latent_width, + block_state.clip_len, + device=device, + ).to(block_state.condition_latents.dtype) + block_state.condition_y = torch.cat([condition_mask, block_state.condition_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 `y`: the previous " + "segment's tail frames (zeros for the first segment) are VAE-encoded, masked, and stacked under the " + "reference half `y_ref`. 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("y_ref", required=True, type_hint=torch.Tensor), + InputParam("clip_len", type_hint=int, default=81), + InputParam("first_num", type_hint=int, default=1), + InputParam("height", required=True, type_hint=int), + InputParam("width", required=True, type_hint=int), + InputParam("latent_height", required=True, type_hint=int), + InputParam("latent_width", required=True, type_hint=int), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + "y", + 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 + + num_frames = block_state.clip_len + 1 + mask_len = block_state.first_num 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=(block_state.height, block_state.width), mode="bicubic" + ).permute(1, 0, 2, 3) + cond_pixels = torch.cat( + [ + prev_frames, + torch.zeros(3, num_frames - mask_len - 1, block_state.height, block_state.width, device=device), + ], + dim=1, + ) + else: + cond_pixels = torch.zeros(3, num_frames - 1, block_state.height, block_state.width, device=device) + + y_reft = encode_vae(components.vae, cond_pixels.unsqueeze(0)).squeeze(0) + mask_reft = get_i2v_mask( + y_reft.shape[1], block_state.latent_height, block_state.latent_width, mask_len, device=device + ).to(y_reft.dtype) + y_reft = torch.cat([mask_reft, y_reft], dim=0) + + block_state.y = torch.cat([block_state.y_ref, y_reft], 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("generator"), + InputParam("y", required=True, type_hint=torch.Tensor), + InputParam("latent_height", required=True, type_hint=int), + InputParam("latent_width", required=True, type_hint=int), + ] + + @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 = torch.randn( + components.num_channels_latents, + block_state.y.shape[1], + block_state.latent_height, + block_state.latent_width, + dtype=torch.float32, + device=device, + generator=block_state.generator, + ) + 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("num_inference_steps", type_hint=int, 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("condition_latents", required=True, type_hint=torch.Tensor), + InputParam("condition_y", required=True, type_hint=torch.Tensor), + InputParam("condition_clip_context", required=True, type_hint=torch.Tensor), + InputParam("prompt_ref_embeds", required=True, type_hint=torch.Tensor), + InputParam("kv_cache", required=True, type_hint=WanAnimate2KVCache), + InputParam("timesteps", required=True, type_hint=torch.Tensor), + InputParam("max_seq_len_ref", required=True, type_hint=int), + InputParam("grid_sizes_ref", required=True, type_hint=torch.Tensor), + ] + + @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.condition_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.condition_y.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), + InputParam("y", required=True, type_hint=torch.Tensor), + InputParam("kv_cache", required=True, type_hint=WanAnimate2KVCache), + InputParam("timesteps", required=True, type_hint=torch.Tensor), + InputParam("num_inference_steps", type_hint=int, default=40), + InputParam("num_segments", required=True, type_hint=int), + InputParam("max_seq_len", required=True, type_hint=int), + InputParam("grid_sizes_ref", required=True, type_hint=torch.Tensor), + InputParam("clip_len", type_hint=int, default=81), + InputParam("height", required=True, type_hint=int), + InputParam("width", required=True, type_hint=int), + InputParam("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.y.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.clip_len, + 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): + @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), + InputParam("kv_cache", required=True, type_hint=WanAnimate2KVCache), + InputParam("first_num", type_hint=int, default=1), + ] + + @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.first_num :] + + 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), + ] + + @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) + + 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): + 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..ee4565813bab --- /dev/null +++ b/src/diffusers/modular_pipelines/wan_animate_2/encoders.py @@ -0,0 +1,571 @@ +# 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 ...guiders import ClassifierFreeGuidance +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 + + +# Copied from diffusers.pipelines.wan.pipeline_wan_animate_2.clip_visual_encode +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] + + +# Copied from diffusers.pipelines.wan.pipeline_wan_animate_2.get_i2v_mask +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 + + +# Copied from diffusers.pipelines.wan.pipeline_wan_animate_2.get_frame_indices +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 guider " + "needs unconditional embeddings), 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), + ComponentSpec( + "guider", + ClassifierFreeGuidance, + config=FrozenDict({"guidance_scale": 3.0}), + default_creation_method="from_config", + ), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam("prompt", required=True, type_hint=str), + InputParam("negative_prompt", type_hint=str), + InputParam( + "prompt_ref", + default="人物动作的参考视频", + type_hint=str, + description="The reference prompt for the driving video context", + ), + InputParam("max_sequence_length", default=512), + ] + + @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, + ) + block_state.negative_prompt_embeds = None + if components.requires_unconditional_embeds: + 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 WanAnimate2ImageResizeStep(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("image", type_hint=PIL.Image.Image, required=True), + 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_top", type_hint=int, description="Top edge of the content inside the letterbox"), + OutputParam("crop_left", type_hint=int, description="Left edge of the content inside the letterbox"), + OutputParam("crop_height", type_hint=int, description="Height of the content inside the letterbox"), + OutputParam("crop_width", type_hint=int, description="Width of the content inside the letterbox"), + ] + + @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 + block_state.crop_width = ( + width if width / height < image_width / image_height else image_width * height // image_height + ) + block_state.crop_height = ( + height if width / height >= image_width / image_height else image_height * width // image_width + ) + block_state.crop_top = (height - block_state.crop_height) // 2 + block_state.crop_left = (width - block_state.crop_width) // 2 + + 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 WanAnimate2VideoPreprocessStep(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`. Overwritten with the preprocessed `[1, 3, T, H, W]` tensor.", + ), + 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( + "clip_len", type_hint=int, default=81, description="The number of frames in each inference segment" + ), + InputParam( + "first_num", + type_hint=int, + default=1, + description="The number of conditioning frames carried over from the previous segment", + ), + InputParam("height", type_hint=int, required=True), + InputParam("width", type_hint=int, required=True), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + 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 (`clip_len - first_num`)", + ), + ] + + @torch.no_grad() + def __call__(self, components, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + + device = components._execution_device + + 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] + + 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) + + real_frame_len = driving_video.shape[2] + effective_segment = block_state.clip_len - block_state.first_num + if real_frame_len > block_state.first_num: + last_segment_frames = (real_frame_len - block_state.first_num) % effective_segment + else: + last_segment_frames = 0 + num_padding = effective_segment - last_segment_frames if last_segment_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 = 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.first_num + effective_segment - 1 + ) // effective_segment + + self.set_block_state(state, block_state) + return components, state + + +# ======================================== +# Image Encoders (CLIP) +# ======================================== + + +class WanAnimate2ImageEncoderStep(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), + ] + + @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 WanAnimate2DrivingImageEncoderStep(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", + required=True, + type_hint=torch.Tensor, + description="The preprocessed driving video `[1, 3, T, H, W]`", + ), + ] + + @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[0, :, 0], device, components.image_encoder.dtype + ) + + self.set_block_state(state, block_state) + return components, state + + +# ======================================== +# VAE Encoder (reference image) +# ======================================== + + +class WanAnimate2RefVaeEncoderStep(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 `y`" + ) + + @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), + InputParam("height", type_hint=int, required=True), + InputParam("width", type_hint=int, required=True), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + "y_ref", + type_hint=torch.Tensor, + description="i2v mask + reference image latents, `[20, 1, latent_height, latent_width]`", + ), + OutputParam("latent_height", type_hint=int), + OutputParam("latent_width", type_hint=int), + ] + + @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)) + + block_state.latent_height = block_state.height // components.vae_scale_factor_spatial + block_state.latent_width = block_state.width // components.vae_scale_factor_spatial + + mask_ref = get_i2v_mask(1, block_state.latent_height, block_state.latent_width, 1, device=device).to( + ref_latents.dtype + ) + block_state.y_ref = 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..aa112ebc95a2 --- /dev/null +++ b/src/diffusers/modular_pipelines/wan_animate_2/modular_blocks_wan_animate_2.py @@ -0,0 +1,125 @@ +# 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 ...utils import logging +from ..modular_pipeline import SequentialPipelineBlocks +from ..modular_pipeline_utils import OutputParam +from .before_denoise import WanAnimate2PrepareSegmentsStep +from .decoders import WanAnimate2DecodeStep +from .denoise import WanAnimate2DenoiseStep +from .encoders import ( + WanAnimate2DrivingImageEncoderStep, + WanAnimate2ImageEncoderStep, + WanAnimate2ImageResizeStep, + WanAnimate2RefVaeEncoderStep, + WanAnimate2TextEncoderStep, + WanAnimate2VideoPreprocessStep, +) + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +# 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`) + guider (`ClassifierFreeGuidance`) + image_processor (`WanAnimate2VideoProcessor`) + video_processor (`WanAnimate2VideoProcessor`) + image_encoder (`CLIPVisionModel`) + vae (`AutoencoderKLWan`) + transformer (`WanAnimate2Transformer3DModel`) + scheduler (`SchedulerMixin`) + + Inputs: + prompt (`str`): + TODO: Add description. + negative_prompt (`str`, *optional*): + TODO: Add description. + prompt_ref (`str`, *optional*, defaults to 人物动作的参考视频): + The reference prompt for the driving video context + max_sequence_length (`None`, *optional*, defaults to 512): + TODO: Add description. + image (`Image`): + TODO: Add description. + 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`. + Overwritten with the preprocessed `[1, 3, T, H, W]` tensor. + 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 + clip_len (`int`, *optional*, defaults to 81): + The number of frames in each inference segment + first_num (`int`, *optional*, defaults to 1): + The number of conditioning frames carried over from the previous segment + generator (`None`, *optional*): + TODO: Add description. + num_inference_steps (`int`, *optional*, defaults to 40): + TODO: Add description. + **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 = [ + WanAnimate2TextEncoderStep, + WanAnimate2ImageResizeStep, + WanAnimate2VideoPreprocessStep, + WanAnimate2ImageEncoderStep, + WanAnimate2DrivingImageEncoderStep, + WanAnimate2RefVaeEncoderStep, + WanAnimate2PrepareSegmentsStep, + WanAnimate2DenoiseStep, + WanAnimate2DecodeStep, + ] + block_names = [ + "text_encoder", + "image_resize", + "video_preprocess", + "image_encoder", + "driving_image_encoder", + "ref_vae_encoder", + "prepare_segments", + "denoise", + "decode", + ] + + @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..c84e939df70a --- /dev/null +++ b/src/diffusers/modular_pipelines/wan_animate_2/modular_blocks_wan_animate_2_distilled.py @@ -0,0 +1,125 @@ +# 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 ...utils import logging +from ..modular_pipeline import SequentialPipelineBlocks +from ..modular_pipeline_utils import OutputParam +from .before_denoise import WanAnimate2PrepareSegmentsStep +from .decoders import WanAnimate2DecodeStep +from .denoise import WanAnimate2DistilledDenoiseStep +from .encoders import ( + WanAnimate2DrivingImageEncoderStep, + WanAnimate2ImageEncoderStep, + WanAnimate2ImageResizeStep, + WanAnimate2RefVaeEncoderStep, + WanAnimate2TextEncoderStep, + WanAnimate2VideoPreprocessStep, +) + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +# 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`) + guider (`ClassifierFreeGuidance`) + image_processor (`WanAnimate2VideoProcessor`) + video_processor (`WanAnimate2VideoProcessor`) + image_encoder (`CLIPVisionModel`) + vae (`AutoencoderKLWan`) + transformer (`WanAnimate2Transformer3DModel`) + scheduler (`SchedulerMixin`) + + Inputs: + prompt (`str`): + TODO: Add description. + negative_prompt (`str`, *optional*): + TODO: Add description. + prompt_ref (`str`, *optional*, defaults to 人物动作的参考视频): + The reference prompt for the driving video context + max_sequence_length (`None`, *optional*, defaults to 512): + TODO: Add description. + image (`Image`): + TODO: Add description. + 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`. + Overwritten with the preprocessed `[1, 3, T, H, W]` tensor. + 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 + clip_len (`int`, *optional*, defaults to 81): + The number of frames in each inference segment + first_num (`int`, *optional*, defaults to 1): + The number of conditioning frames carried over from the previous segment + generator (`None`, *optional*): + TODO: Add description. + num_inference_steps (`int`, *optional*, defaults to 40): + TODO: Add description. + **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 = [ + WanAnimate2TextEncoderStep, + WanAnimate2ImageResizeStep, + WanAnimate2VideoPreprocessStep, + WanAnimate2ImageEncoderStep, + WanAnimate2DrivingImageEncoderStep, + WanAnimate2RefVaeEncoderStep, + WanAnimate2PrepareSegmentsStep, + WanAnimate2DistilledDenoiseStep, + WanAnimate2DecodeStep, + ] + block_names = [ + "text_encoder", + "image_resize", + "video_preprocess", + "image_encoder", + "driving_image_encoder", + "ref_vae_encoder", + "prepare_segments", + "denoise", + "decode", + ] + + @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..742a55147063 --- /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/pipelines/wan/image_processor.py b/src/diffusers/pipelines/wan/image_processor.py index f1a046a515c6..fa18150fcc6e 100644 --- a/src/diffusers/pipelines/wan/image_processor.py +++ b/src/diffusers/pipelines/wan/image_processor.py @@ -20,7 +20,6 @@ from ...configuration_utils import register_to_config from ...image_processor import VaeImageProcessor from ...utils import PIL_INTERPOLATION -from ...video_processor import VideoProcessor class WanAnimateImageProcessor(VaeImageProcessor): @@ -183,58 +182,3 @@ def get_default_height_width( width = round(np.sqrt(max_area / aspect_ratio)) // mod_value_w * mod_value_w return height, width - - -class WanAnimate2VideoProcessor(VideoProcessor, WanAnimateImageProcessor): - 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). Same letterbox as - [`WanAnimateImageProcessor`], except the resized content is pasted at - `((height - src_h) // 2, (width - src_w) // 2)` -- the placement convention of the reference implementation -- - instead of `(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__`s: they are themselves - # `register_to_config`-decorated, so calling one 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 parents' bodies only validate. - 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 diff --git a/src/diffusers/pipelines/wan/pipeline_wan_animate_2.py b/src/diffusers/pipelines/wan/pipeline_wan_animate_2.py index 834afb2afc38..bcc6dded0da3 100644 --- a/src/diffusers/pipelines/wan/pipeline_wan_animate_2.py +++ b/src/diffusers/pipelines/wan/pipeline_wan_animate_2.py @@ -24,10 +24,10 @@ from ...loaders import WanLoraLoaderMixin from ...models import AutoencoderKLWan, WanAnimate2Transformer3DModel from ...models.transformers.transformer_wan_animate_2 import WanAnimate2KVCache +from ...modular_pipelines.wan_animate_2.video_processor import WanAnimate2VideoProcessor from ...schedulers import SchedulerMixin from ...utils import logging from ..pipeline_utils import DiffusionPipeline -from .image_processor import WanAnimate2VideoProcessor from .pipeline_output import WanPipelineOutput diff --git a/src/diffusers/utils/dummy_torch_and_transformers_objects.py b/src/diffusers/utils/dummy_torch_and_transformers_objects.py index ce6fc692a16d..cf348392ad0d 100644 --- a/src/diffusers/utils/dummy_torch_and_transformers_objects.py +++ b/src/diffusers/utils/dummy_torch_and_transformers_objects.py @@ -782,6 +782,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"] From cc5c239a4a7bc45f7caf6df36e791a60aaa4c2e0 Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Tue, 11 Aug 2026 06:19:51 +0000 Subject: [PATCH 08/19] Repack Wan-Animate-2 modular blocks into canonical standalone steps - Top-level blockset children now follow the family convention (text_encoder / image_encoder / video_encoder / vae_encoder / denoise / decode), each poppable and usable standalone, each a flat sequence of leaf blocks (wan-i2v / flux2 shape). - Hoist the driving-video VAE encode out of the segment loop into the vae_encoder group: the Wan VAE is causal in time so each slice is encoded separately, but all slices are known upfront. The loop keeps only the genuinely sequential work (prev-frame conditioning and per-segment decode). - Drop the guider from the text encoder step: the denoise step owns the guider spec (removing the 3.0-vs-1.0 spec conflict in the distilled preset), and the text encoder encodes the negative prompt when the pipeline's guider requires unconditional embeddings or one is passed explicitly -- standalone, nothing is encoded unless asked. - Rename to canonical step names: ProcessImagesInputStep / ProcessVideosInputStep (flux/qwen convention), {Image,Video}{Clip,Vae} EncoderStep leaves, short EncodeStep group names. - Rename clip_len -> segment_frame_length and first_num -> prev_segment_conditioning_frames, matching the merged Wan-Animate v1 pipeline's argument names. Verified bit-identical to the standard pipeline after each change (tiny fp32 base+distilled, real distilled checkpoint bf16). Co-Authored-By: Claude Fable 5 --- .../wan_animate_2/before_denoise.py | 6 +- .../wan_animate_2/denoise.py | 112 ++---- .../wan_animate_2/encoders.py | 131 +++++-- .../modular_blocks_wan_animate_2.py | 332 ++++++++++++++++-- .../modular_blocks_wan_animate_2_distilled.py | 147 ++++++-- .../pipelines/wan/pipeline_wan_animate_2.py | 50 +-- 6 files changed, 578 insertions(+), 200 deletions(-) diff --git a/src/diffusers/modular_pipelines/wan_animate_2/before_denoise.py b/src/diffusers/modular_pipelines/wan_animate_2/before_denoise.py index 25dd43ddc0d2..583899c87f8f 100644 --- a/src/diffusers/modular_pipelines/wan_animate_2/before_denoise.py +++ b/src/diffusers/modular_pipelines/wan_animate_2/before_denoise.py @@ -32,14 +32,14 @@ class WanAnimate2PrepareSegmentsStep(ModularPipelineBlocks): def description(self) -> str: return ( "Step that computes the segment-invariant geometry for the segment loop. The zigzag padding makes " - "every segment exactly `clip_len` frames, so the latent grid, the packed sequence lengths, and the " + "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("clip_len", type_hint=int, default=81), + InputParam("segment_frame_length", type_hint=int, default=81), InputParam("latent_height", type_hint=int, required=True), InputParam("latent_width", type_hint=int, required=True), ] @@ -69,7 +69,7 @@ def intermediate_outputs(self) -> list[OutputParam]: def __call__(self, components, state: PipelineState) -> PipelineState: block_state = self.get_block_state(state) - latent_segment_frames = (block_state.clip_len - 1) // components.vae_scale_factor_temporal + 1 + latent_segment_frames = (block_state.segment_frame_length - 1) // components.vae_scale_factor_temporal + 1 ref_shape = [latent_segment_frames, block_state.latent_height, block_state.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) diff --git a/src/diffusers/modular_pipelines/wan_animate_2/denoise.py b/src/diffusers/modular_pipelines/wan_animate_2/denoise.py index eac81a5cf7a6..aa845be8815a 100644 --- a/src/diffusers/modular_pipelines/wan_animate_2/denoise.py +++ b/src/diffusers/modular_pipelines/wan_animate_2/denoise.py @@ -50,71 +50,6 @@ def decode_vae(vae: AutoencoderKLWan, latents: torch.Tensor) -> torch.Tensor: # ======================================== -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. The Wan VAE " - "is causal in time, so encoding the whole video once up front and slicing the latents would not be " - "equivalent — each segment restarts the temporal convolution. 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", required=True, type_hint=torch.Tensor), - InputParam("effective_segment", required=True, type_hint=int), - InputParam("clip_len", type_hint=int, default=81), - InputParam("latent_height", required=True, type_hint=int), - InputParam("latent_width", required=True, type_hint=int), - ] - - @property - def intermediate_outputs(self) -> list[OutputParam]: - return [ - OutputParam( - "condition_latents", - type_hint=torch.Tensor, - description="VAE latents of this segment's driving-video slice", - ), - OutputParam( - "condition_y", - 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 - - start = k * block_state.effective_segment - block_state.condition_latents = encode_vae( - components.vae, block_state.driving_video[:, :, start : start + block_state.clip_len] - ) - - condition_mask = get_i2v_mask( - block_state.condition_latents.shape[2], - block_state.latent_height, - block_state.latent_width, - block_state.clip_len, - device=device, - ).to(block_state.condition_latents.dtype) - block_state.condition_y = torch.cat([condition_mask, block_state.condition_latents[0]], dim=0) - - return components, block_state - - class WanAnimate2SegmentPrevFramesStep(ModularPipelineBlocks): model_name = "wan-animate-2" @@ -138,8 +73,8 @@ def expected_components(self) -> list[ComponentSpec]: def inputs(self) -> list[InputParam]: return [ InputParam("y_ref", required=True, type_hint=torch.Tensor), - InputParam("clip_len", type_hint=int, default=81), - InputParam("first_num", type_hint=int, default=1), + InputParam("segment_frame_length", type_hint=int, default=81), + InputParam("prev_segment_conditioning_frames", type_hint=int, default=1), InputParam("height", required=True, type_hint=int), InputParam("width", required=True, type_hint=int), InputParam("latent_height", required=True, type_hint=int), @@ -162,8 +97,8 @@ def __call__(self, components, block_state: BlockState, k: int): # previous iteration. device = components._execution_device - num_frames = block_state.clip_len + 1 - mask_len = block_state.first_num if k > 0 else 0 + 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( @@ -305,8 +240,18 @@ def expected_components(self) -> list[ComponentSpec]: @property def inputs(self) -> list[InputParam]: return [ - InputParam("condition_latents", required=True, type_hint=torch.Tensor), - InputParam("condition_y", required=True, type_hint=torch.Tensor), + InputParam( + "condition_latents", + required=True, + type_hint=torch.Tensor, + description="VAE latents of every segment's driving-video slice, `[num_segments, 16, T, H', W']`", + ), + InputParam( + "condition_y", + required=True, + type_hint=torch.Tensor, + description="i2v mask + driving-slice latents per segment, `[num_segments, 20, T, H', W']`", + ), InputParam("condition_clip_context", required=True, type_hint=torch.Tensor), InputParam("prompt_ref_embeds", required=True, type_hint=torch.Tensor), InputParam("kv_cache", required=True, type_hint=WanAnimate2KVCache), @@ -322,11 +267,11 @@ def __call__(self, components, block_state: BlockState, k: int): t_ref = torch.tensor([block_state.timesteps[0].item()], device=device, dtype=transformer_dtype) components.transformer( - [block_state.condition_latents[0].to(transformer_dtype)], + [block_state.condition_latents[k].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.condition_y.to(transformer_dtype)], + condition_latents=[block_state.condition_y[k].to(transformer_dtype)], kv_cache=block_state.kv_cache, kv_cache_mode="extract", seq_len=block_state.max_seq_len_ref, @@ -377,7 +322,7 @@ def inputs(self) -> list[InputParam]: InputParam("num_segments", required=True, type_hint=int), InputParam("max_seq_len", required=True, type_hint=int), InputParam("grid_sizes_ref", required=True, type_hint=torch.Tensor), - InputParam("clip_len", type_hint=int, default=81), + InputParam("segment_frame_length", type_hint=int, default=81), InputParam("height", required=True, type_hint=int), InputParam("width", required=True, type_hint=int), InputParam("generator"), @@ -431,7 +376,7 @@ def __call__(self, components, block_state: BlockState, k: int): kv_cache_mode="cached", seq_len=block_state.max_seq_len, reference_grid_sizes=block_state.grid_sizes_ref, - origin_len=block_state.clip_len, + origin_len=block_state.segment_frame_length, origin_area=[block_state.height, block_state.width], is_uncondtion=guider_state_batch.is_uncondtion, **shared_kwargs, @@ -508,7 +453,7 @@ def inputs(self) -> list[InputParam]: return [ InputParam("latents", required=True, type_hint=torch.Tensor), InputParam("kv_cache", required=True, type_hint=WanAnimate2KVCache), - InputParam("first_num", type_hint=int, default=1), + InputParam("prev_segment_conditioning_frames", type_hint=int, default=1), ] @property @@ -528,7 +473,7 @@ def __call__(self, components, block_state: BlockState, k: int): out_frames = decode_vae(components.vae, latents[:, 1:]) if k > 0: - out_frames = out_frames[:, :, block_state.first_num :] + out_frames = out_frames[:, :, block_state.prev_segment_conditioning_frames :] block_state.segment_frames.append(out_frames.cpu()) block_state.out_frames = out_frames @@ -553,8 +498,8 @@ class WanAnimate2SegmentLoopWrapper(LoopSequentialPipelineBlocks): 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." + "for preparation, reference extraction, denoising, and decoding; each segment conditions on the " + "previous one's decoded tail frames." ) @property @@ -594,7 +539,6 @@ def __call__(self, components, state: PipelineState) -> PipelineState: class WanAnimate2DenoiseStep(WanAnimate2SegmentLoopWrapper): block_classes = [ - WanAnimate2SegmentVaeEncoderStep, WanAnimate2SegmentPrevFramesStep, WanAnimate2SegmentPrepareStep, WanAnimate2SegmentSchedulerResetStep, @@ -603,7 +547,6 @@ class WanAnimate2DenoiseStep(WanAnimate2SegmentLoopWrapper): WanAnimate2SegmentDecodeStep, ] block_names = [ - "vae_encoder", "prev_frames", "prepare", "scheduler_reset", @@ -616,14 +559,12 @@ class WanAnimate2DenoiseStep(WanAnimate2SegmentLoopWrapper): 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." + "At each segment: prev_frames -> prepare -> scheduler_reset -> ref_extract -> denoise_inner -> decode." ) class WanAnimate2DistilledDenoiseStep(WanAnimate2SegmentLoopWrapper): block_classes = [ - WanAnimate2SegmentVaeEncoderStep, WanAnimate2SegmentPrevFramesStep, WanAnimate2SegmentPrepareStep, WanAnimate2SegmentSchedulerResetStep, @@ -632,7 +573,6 @@ class WanAnimate2DistilledDenoiseStep(WanAnimate2SegmentLoopWrapper): WanAnimate2SegmentDecodeStep, ] block_names = [ - "vae_encoder", "prev_frames", "prepare", "scheduler_reset", @@ -645,6 +585,6 @@ class WanAnimate2DistilledDenoiseStep(WanAnimate2SegmentLoopWrapper): 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 -> " + "At each segment: 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 index ee4565813bab..907f1de0bfc5 100644 --- a/src/diffusers/modular_pipelines/wan_animate_2/encoders.py +++ b/src/diffusers/modular_pipelines/wan_animate_2/encoders.py @@ -21,7 +21,6 @@ from transformers import AutoTokenizer, CLIPVisionModel, UMT5EncoderModel from ...configuration_utils import FrozenDict -from ...guiders import ClassifierFreeGuidance from ...models import AutoencoderKLWan from ...utils import logging from ..modular_pipeline import ModularPipelineBlocks, PipelineState @@ -129,8 +128,9 @@ class WanAnimate2TextEncoderStep(ModularPipelineBlocks): @property def description(self) -> str: return ( - "Text Encoder step that encodes the character/background prompt, the negative prompt (when the guider " - "needs unconditional embeddings), and the fixed reference prompt for the driving-video context" + "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 @@ -138,12 +138,6 @@ def expected_components(self) -> list[ComponentSpec]: return [ ComponentSpec("text_encoder", UMT5EncoderModel), ComponentSpec("tokenizer", AutoTokenizer), - ComponentSpec( - "guider", - ClassifierFreeGuidance, - config=FrozenDict({"guidance_scale": 3.0}), - default_creation_method="from_config", - ), ] @property @@ -191,8 +185,12 @@ def __call__(self, components, state: PipelineState) -> PipelineState: 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: + 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, @@ -217,7 +215,7 @@ def __call__(self, components, state: PipelineState) -> PipelineState: # ======================================== -class WanAnimate2ImageResizeStep(ModularPipelineBlocks): +class WanAnimate2ProcessImagesInputStep(ModularPipelineBlocks): model_name = "wan-animate-2" @property @@ -303,7 +301,7 @@ def __call__(self, components, state: PipelineState) -> PipelineState: return components, state -class WanAnimate2VideoPreprocessStep(ModularPipelineBlocks): +class WanAnimate2ProcessVideosInputStep(ModularPipelineBlocks): model_name = "wan-animate-2" @property @@ -345,10 +343,13 @@ def inputs(self) -> list[InputParam]: ), InputParam("fps", type_hint=int, default=24, description="The frame rate the model generates at"), InputParam( - "clip_len", type_hint=int, default=81, description="The number of frames in each inference segment" + "segment_frame_length", + type_hint=int, + default=81, + description="The number of frames in each inference segment", ), InputParam( - "first_num", + "prev_segment_conditioning_frames", type_hint=int, default=1, description="The number of conditioning frames carried over from the previous segment", @@ -369,7 +370,7 @@ def intermediate_outputs(self) -> list[OutputParam]: OutputParam( "effective_segment", type_hint=int, - description="Frames each segment advances by (`clip_len - first_num`)", + description="Frames each segment advances by (`segment_frame_length - prev_segment_conditioning_frames`)", ), ] @@ -389,9 +390,9 @@ def __call__(self, components, state: PipelineState) -> PipelineState: ).to(device, dtype=torch.float32) real_frame_len = driving_video.shape[2] - effective_segment = block_state.clip_len - block_state.first_num - if real_frame_len > block_state.first_num: - last_segment_frames = (real_frame_len - block_state.first_num) % effective_segment + effective_segment = block_state.segment_frame_length - block_state.prev_segment_conditioning_frames + if real_frame_len > block_state.prev_segment_conditioning_frames: + last_segment_frames = (real_frame_len - block_state.prev_segment_conditioning_frames) % effective_segment else: last_segment_frames = 0 num_padding = effective_segment - last_segment_frames if last_segment_frames > 0 else 0 @@ -405,7 +406,7 @@ def __call__(self, components, state: PipelineState) -> PipelineState: block_state.real_frame_len = real_frame_len block_state.effective_segment = effective_segment block_state.num_segments = ( - target_num_frames - block_state.first_num + effective_segment - 1 + target_num_frames - block_state.prev_segment_conditioning_frames + effective_segment - 1 ) // effective_segment self.set_block_state(state, block_state) @@ -417,7 +418,7 @@ def __call__(self, components, state: PipelineState) -> PipelineState: # ======================================== -class WanAnimate2ImageEncoderStep(ModularPipelineBlocks): +class WanAnimate2ImageClipEncoderStep(ModularPipelineBlocks): model_name = "wan-animate-2" @property @@ -460,7 +461,7 @@ def __call__(self, components, state: PipelineState) -> PipelineState: return components, state -class WanAnimate2DrivingImageEncoderStep(ModularPipelineBlocks): +class WanAnimate2VideoClipEncoderStep(ModularPipelineBlocks): model_name = "wan-animate-2" @property @@ -511,11 +512,11 @@ def __call__(self, components, state: PipelineState) -> PipelineState: # ======================================== -# VAE Encoder (reference image) +# VAE Encoders # ======================================== -class WanAnimate2RefVaeEncoderStep(ModularPipelineBlocks): +class WanAnimate2ImageVaeEncoderStep(ModularPipelineBlocks): model_name = "wan-animate-2" @property @@ -569,3 +570,87 @@ def __call__(self, components, state: PipelineState) -> PipelineState: self.set_block_state(state, block_state) return components, state + + +class WanAnimate2VideoVaeEncoderStep(ModularPipelineBlocks): + model_name = "wan-animate-2" + + @property + def description(self) -> str: + return ( + "VAE Encoder step that encodes every 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." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("vae", AutoencoderKLWan), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + "driving_video", + required=True, + type_hint=torch.Tensor, + description="The preprocessed driving video `[1, 3, T, H, W]`", + ), + InputParam("num_segments", required=True, type_hint=int), + InputParam("effective_segment", required=True, type_hint=int), + InputParam("segment_frame_length", type_hint=int, default=81), + InputParam("latent_height", required=True, type_hint=int), + InputParam("latent_width", required=True, type_hint=int), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + "condition_latents", + type_hint=torch.Tensor, + description="VAE latents of every segment's driving-video slice, `[num_segments, 16, T, H', W']`", + ), + OutputParam( + "condition_y", + type_hint=torch.Tensor, + description=( + "i2v mask + driving-slice latents per segment, `[num_segments, 20, T, H', W']`, conditioning " + "the reference-extraction pass" + ), + ), + ] + + @torch.no_grad() + def __call__(self, components, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + + device = components._execution_device + + condition_latents = [] + for k in range(block_state.num_segments): + start = k * block_state.effective_segment + condition_latents.append( + encode_vae( + components.vae, block_state.driving_video[:, :, start : start + block_state.segment_frame_length] + ) + ) + block_state.condition_latents = torch.cat(condition_latents, dim=0) + + # After zigzag padding every segment is exactly `segment_frame_length` frames, so one mask fits all. + condition_mask = get_i2v_mask( + block_state.condition_latents.shape[2], + block_state.latent_height, + block_state.latent_width, + block_state.segment_frame_length, + device=device, + ).to(block_state.condition_latents.dtype) + block_state.condition_y = torch.stack( + [torch.cat([condition_mask, latents], dim=0) for latents in block_state.condition_latents] + ) + + 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 index aa112ebc95a2..b58a474acde1 100644 --- 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 @@ -14,23 +14,311 @@ from ...utils import logging from ..modular_pipeline import SequentialPipelineBlocks -from ..modular_pipeline_utils import OutputParam +from ..modular_pipeline_utils import InsertableDict, OutputParam from .before_denoise import WanAnimate2PrepareSegmentsStep from .decoders import WanAnimate2DecodeStep from .denoise import WanAnimate2DenoiseStep from .encoders import ( - WanAnimate2DrivingImageEncoderStep, - WanAnimate2ImageEncoderStep, - WanAnimate2ImageResizeStep, - WanAnimate2RefVaeEncoderStep, + WanAnimate2ImageClipEncoderStep, + WanAnimate2ImageVaeEncoderStep, + WanAnimate2ProcessImagesInputStep, + WanAnimate2ProcessVideosInputStep, WanAnimate2TextEncoderStep, - WanAnimate2VideoPreprocessStep, + WanAnimate2VideoClipEncoderStep, + WanAnimate2VideoVaeEncoderStep, ) 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`): + TODO: Add description. + 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_top (`int`): + Top edge of the content inside the letterbox + crop_left (`int`): + Left edge of the content inside the letterbox + crop_height (`int`): + Height of the content inside the letterbox + crop_width (`int`): + Width of the content inside the letterbox + 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`. + Overwritten with the preprocessed `[1, 3, T, H, W]` tensor. + 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`): + TODO: Add description. + width (`int`): + TODO: Add description. + + Outputs: + 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`." + ) + + +WanAnimate2VaeEncoderBlocks = InsertableDict( + [ + ("image_encode", WanAnimate2ImageVaeEncoderStep()), + ("video_encode", WanAnimate2VideoVaeEncoderStep()), + ] +) + + +# auto_docstring +class WanAnimate2VaeEncodeStep(SequentialPipelineBlocks): + """ + VAE encoder step that encodes the letterboxed reference image into the reference half of the conditioning tensor `y_ref`, and every segment's slice of the driving video into the reference-extraction conditioning `condition_latents` / `condition_y`. + + Components: + vae (`AutoencoderKLWan`) + + Inputs: + image_pixels (`Tensor`): + TODO: Add description. + height (`int`): + TODO: Add description. + width (`int`): + TODO: Add description. + driving_video (`Tensor`): + The preprocessed driving video `[1, 3, T, H, W]` + num_segments (`int`): + TODO: Add description. + effective_segment (`int`): + TODO: Add description. + segment_frame_length (`int`, *optional*, defaults to 81): + TODO: Add description. + + Outputs: + y_ref (`Tensor`): + i2v mask + reference image latents, `[20, 1, latent_height, latent_width]` + latent_height (`int`): + TODO: Add description. + latent_width (`int`): + TODO: Add description. + condition_latents (`Tensor`): + VAE latents of every segment's driving-video slice, `[num_segments, 16, T, H', W']` + condition_y (`Tensor`): + i2v mask + driving-slice latents per segment, `[num_segments, 20, T, H', W']`, conditioning the + reference-extraction pass + """ + + model_name = "wan-animate-2" + block_classes = WanAnimate2VaeEncoderBlocks.values() + block_names = WanAnimate2VaeEncoderBlocks.keys() + + @property + def description(self): + return ( + "VAE encoder step that encodes the letterboxed reference image into the reference half of the " + "conditioning tensor `y_ref`, and every segment's slice of the driving video into the " + "reference-extraction conditioning `condition_latents` / `condition_y`." + ) + + +# ==================== +# 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): + TODO: Add description. + latent_height (`int`): + TODO: Add description. + latent_width (`int`): + TODO: Add description. + num_segments (`int`): + TODO: Add description. + y_ref (`Tensor`): + TODO: Add description. + prev_segment_conditioning_frames (`int`, *optional*, defaults to 1): + TODO: Add description. + height (`int`): + TODO: Add description. + width (`int`): + TODO: Add description. + generator (`None`, *optional*): + TODO: Add description. + num_inference_steps (`int`, *optional*, defaults to 40): + TODO: Add description. + condition_latents (`Tensor`): + VAE latents of every segment's driving-video slice, `[num_segments, 16, T, H', W']` + condition_y (`Tensor`): + i2v mask + driving-slice latents per segment, `[num_segments, 20, T, H', W']` + condition_clip_context (`Tensor`): + TODO: Add description. + prompt_ref_embeds (`Tensor`): + TODO: Add description. + 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: + grid_sizes_ref (`Tensor`): + 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 + max_seq_len (`int`): + Packed sequence length of the generation tokens + max_seq_len_ref (`int`): + Packed sequence length of the reference tokens + y (`Tensor`): + The full conditioning tensor: reference half stacked over the segment half + latents (`Tensor`): + This segment's initial noise + kv_cache (`WanAnimate2KVCache`): + Fresh per-segment cache for the reference K/V + timesteps (`Tensor`): + This segment's denoising timesteps + out_frames (`Tensor`): + This segment's decoded frames on device; the next segment conditions on its tail + segment_frames (`list`): + Per-segment decoded frames on CPU, each `[1, 3, T, H, W]` + """ + + 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." + ) + + +# ==================== +# 3. Blocks +# ==================== + + +BLOCKS = InsertableDict( + [ + ("text_encoder", WanAnimate2TextEncoderStep()), + ("image_encoder", WanAnimate2ImageEncodeStep()), + ("video_encoder", WanAnimate2VideoEncodeStep()), + ("vae_encoder", WanAnimate2VaeEncodeStep()), + ("denoise", WanAnimate2CoreDenoiseStep()), + ("decode", WanAnimate2DecodeStep()), + ] +) + + # auto_docstring class WanAnimate2Blocks(SequentialPipelineBlocks): """ @@ -39,13 +327,13 @@ class WanAnimate2Blocks(SequentialPipelineBlocks): Components: text_encoder (`UMT5EncoderModel`) tokenizer (`AutoTokenizer`) - guider (`ClassifierFreeGuidance`) image_processor (`WanAnimate2VideoProcessor`) - video_processor (`WanAnimate2VideoProcessor`) image_encoder (`CLIPVisionModel`) + video_processor (`WanAnimate2VideoProcessor`) vae (`AutoencoderKLWan`) transformer (`WanAnimate2Transformer3DModel`) scheduler (`SchedulerMixin`) + guider (`ClassifierFreeGuidance`) Inputs: prompt (`str`): @@ -71,9 +359,9 @@ class WanAnimate2Blocks(SequentialPipelineBlocks): 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 - clip_len (`int`, *optional*, defaults to 81): + segment_frame_length (`int`, *optional*, defaults to 81): The number of frames in each inference segment - first_num (`int`, *optional*, defaults to 1): + prev_segment_conditioning_frames (`int`, *optional*, defaults to 1): The number of conditioning frames carried over from the previous segment generator (`None`, *optional*): TODO: Add description. @@ -90,28 +378,8 @@ class WanAnimate2Blocks(SequentialPipelineBlocks): """ model_name = "wan-animate-2" - block_classes = [ - WanAnimate2TextEncoderStep, - WanAnimate2ImageResizeStep, - WanAnimate2VideoPreprocessStep, - WanAnimate2ImageEncoderStep, - WanAnimate2DrivingImageEncoderStep, - WanAnimate2RefVaeEncoderStep, - WanAnimate2PrepareSegmentsStep, - WanAnimate2DenoiseStep, - WanAnimate2DecodeStep, - ] - block_names = [ - "text_encoder", - "image_resize", - "video_preprocess", - "image_encoder", - "driving_image_encoder", - "ref_vae_encoder", - "prepare_segments", - "denoise", - "decode", - ] + block_classes = BLOCKS.values() + block_names = BLOCKS.keys() @property def description(self): 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 index c84e939df70a..5394538f3eb8 100644 --- 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 @@ -14,23 +14,122 @@ from ...utils import logging from ..modular_pipeline import SequentialPipelineBlocks -from ..modular_pipeline_utils import OutputParam +from ..modular_pipeline_utils import InsertableDict, OutputParam from .before_denoise import WanAnimate2PrepareSegmentsStep from .decoders import WanAnimate2DecodeStep from .denoise import WanAnimate2DistilledDenoiseStep -from .encoders import ( - WanAnimate2DrivingImageEncoderStep, - WanAnimate2ImageEncoderStep, - WanAnimate2ImageResizeStep, - WanAnimate2RefVaeEncoderStep, - WanAnimate2TextEncoderStep, - WanAnimate2VideoPreprocessStep, +from .encoders import WanAnimate2TextEncoderStep +from .modular_blocks_wan_animate_2 import ( + WanAnimate2ImageEncodeStep, + WanAnimate2VaeEncodeStep, + WanAnimate2VideoEncodeStep, ) logger = logging.get_logger(__name__) # pylint: disable=invalid-name +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): + TODO: Add description. + latent_height (`int`): + TODO: Add description. + latent_width (`int`): + TODO: Add description. + num_segments (`int`): + TODO: Add description. + y_ref (`Tensor`): + TODO: Add description. + prev_segment_conditioning_frames (`int`, *optional*, defaults to 1): + TODO: Add description. + height (`int`): + TODO: Add description. + width (`int`): + TODO: Add description. + generator (`None`, *optional*): + TODO: Add description. + num_inference_steps (`int`, *optional*, defaults to 40): + TODO: Add description. + condition_latents (`Tensor`): + VAE latents of every segment's driving-video slice, `[num_segments, 16, T, H', W']` + condition_y (`Tensor`): + i2v mask + driving-slice latents per segment, `[num_segments, 20, T, H', W']` + condition_clip_context (`Tensor`): + TODO: Add description. + prompt_ref_embeds (`Tensor`): + TODO: Add description. + 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: + grid_sizes_ref (`Tensor`): + 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 + max_seq_len (`int`): + Packed sequence length of the generation tokens + max_seq_len_ref (`int`): + Packed sequence length of the reference tokens + y (`Tensor`): + The full conditioning tensor: reference half stacked over the segment half + latents (`Tensor`): + This segment's initial noise + kv_cache (`WanAnimate2KVCache`): + Fresh per-segment cache for the reference K/V + timesteps (`Tensor`): + This segment's denoising timesteps + out_frames (`Tensor`): + This segment's decoded frames on device; the next segment conditions on its tail + segment_frames (`list`): + Per-segment decoded frames on CPU, each `[1, 3, T, H, W]` + """ + + model_name = "wan-animate-2" + 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." + ) + + +DISTILLED_BLOCKS = InsertableDict( + [ + ("text_encoder", WanAnimate2TextEncoderStep()), + ("image_encoder", WanAnimate2ImageEncodeStep()), + ("video_encoder", WanAnimate2VideoEncodeStep()), + ("vae_encoder", WanAnimate2VaeEncodeStep()), + ("denoise", WanAnimate2DistilledCoreDenoiseStep()), + ("decode", WanAnimate2DecodeStep()), + ] +) + + # auto_docstring class WanAnimate2DistilledBlocks(SequentialPipelineBlocks): """ @@ -39,13 +138,13 @@ class WanAnimate2DistilledBlocks(SequentialPipelineBlocks): Components: text_encoder (`UMT5EncoderModel`) tokenizer (`AutoTokenizer`) - guider (`ClassifierFreeGuidance`) image_processor (`WanAnimate2VideoProcessor`) - video_processor (`WanAnimate2VideoProcessor`) image_encoder (`CLIPVisionModel`) + video_processor (`WanAnimate2VideoProcessor`) vae (`AutoencoderKLWan`) transformer (`WanAnimate2Transformer3DModel`) scheduler (`SchedulerMixin`) + guider (`ClassifierFreeGuidance`) Inputs: prompt (`str`): @@ -71,9 +170,9 @@ class WanAnimate2DistilledBlocks(SequentialPipelineBlocks): 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 - clip_len (`int`, *optional*, defaults to 81): + segment_frame_length (`int`, *optional*, defaults to 81): The number of frames in each inference segment - first_num (`int`, *optional*, defaults to 1): + prev_segment_conditioning_frames (`int`, *optional*, defaults to 1): The number of conditioning frames carried over from the previous segment generator (`None`, *optional*): TODO: Add description. @@ -90,28 +189,8 @@ class WanAnimate2DistilledBlocks(SequentialPipelineBlocks): """ model_name = "wan-animate-2" - block_classes = [ - WanAnimate2TextEncoderStep, - WanAnimate2ImageResizeStep, - WanAnimate2VideoPreprocessStep, - WanAnimate2ImageEncoderStep, - WanAnimate2DrivingImageEncoderStep, - WanAnimate2RefVaeEncoderStep, - WanAnimate2PrepareSegmentsStep, - WanAnimate2DistilledDenoiseStep, - WanAnimate2DecodeStep, - ] - block_names = [ - "text_encoder", - "image_resize", - "video_preprocess", - "image_encoder", - "driving_image_encoder", - "ref_vae_encoder", - "prepare_segments", - "denoise", - "decode", - ] + block_classes = DISTILLED_BLOCKS.values() + block_names = DISTILLED_BLOCKS.keys() @property def description(self): diff --git a/src/diffusers/pipelines/wan/pipeline_wan_animate_2.py b/src/diffusers/pipelines/wan/pipeline_wan_animate_2.py index bcc6dded0da3..888b4191417d 100644 --- a/src/diffusers/pipelines/wan/pipeline_wan_animate_2.py +++ b/src/diffusers/pipelines/wan/pipeline_wan_animate_2.py @@ -240,8 +240,8 @@ def __call__( prompt_ref: str = "人物动作的参考视频", height: int = 800, width: int = 640, - clip_len: int = 81, - first_num: int = 1, + segment_frame_length: int = 81, + prev_segment_conditioning_frames: int = 1, fps: int = 24, driving_video_fps: float | None = None, num_inference_steps: int = 40, @@ -275,9 +275,9 @@ def __call__( rescaled to hit that area and then floored to a multiple of 16. width (`int`, defaults to `640`): See `height`. - clip_len (`int`, defaults to `81`): + segment_frame_length (`int`, defaults to `81`): The number of frames in each inference segment. - first_num (`int`, defaults to `1`): + prev_segment_conditioning_frames (`int`, defaults to `1`): The number of conditioning frames from the previous segment. fps (`int`, defaults to `24`): The frame rate the model generates at. `driving_video` is resampled to it when @@ -335,10 +335,14 @@ def __call__( driving_video, height=actual_h, width=actual_w, resize_mode="fill" ).to(device, dtype=torch.float32) - # Pad driving video to be a multiple of (clip_len - first_num) + # Pad driving video to be a multiple of (segment_frame_length - prev_segment_conditioning_frames) real_frame_len = driving_video.shape[2] - effective_segment = clip_len - first_num - last_segment_frames = (real_frame_len - first_num) % effective_segment if real_frame_len > first_num else 0 + effective_segment = segment_frame_length - prev_segment_conditioning_frames + last_segment_frames = ( + (real_frame_len - prev_segment_conditioning_frames) % effective_segment + if real_frame_len > prev_segment_conditioning_frames + else 0 + ) if last_segment_frames > 0: num_padding = effective_segment - last_segment_frames else: @@ -390,28 +394,30 @@ def __call__( # 7. Segment-based generation loop start = 0 - end = clip_len + end = segment_frame_length all_out_frames = [] out_frames = None - num_segments = (target_num_frames - first_num + effective_segment - 1) // effective_segment + num_segments = ( + target_num_frames - prev_segment_conditioning_frames + effective_segment - 1 + ) // effective_segment for seg_idx in range(num_segments): - if start + first_num >= target_num_frames: + if start + prev_segment_conditioning_frames >= target_num_frames: break - mask_reft_len = first_num if start > 0 else 0 + mask_reft_len = prev_segment_conditioning_frames if start > 0 else 0 - if target_num_frames - start < clip_len: - clip_len_actual = target_num_frames - start + if target_num_frames - start < segment_frame_length: + segment_frame_length_actual = target_num_frames - start else: - clip_len_actual = clip_len + segment_frame_length_actual = segment_frame_length # VAE-encode this segment's slice of the driving video. The Wan VAE is causal in time, so # encoding the whole video once up front and slicing the latents is not the same tensor — - # segments overlap by `first_num` frames and each slice restarts the temporal convolution. + # segments overlap by `prev_segment_conditioning_frames` frames and each slice restarts the temporal convolution. # Encoding per segment is also what a streaming mode would have to do anyway. - condition_latents = self._encode_vae(driving_video[:, :, start : start + clip_len_actual]) + condition_latents = self._encode_vae(driving_video[:, :, start : start + segment_frame_length_actual]) # CLIP features from driving video first frame (direct bicubic to 224×224 from tensor) condition_img = driving_video[0, :, 0] # [C, H, W] in [-1, 1] @@ -420,7 +426,7 @@ def __call__( ) # Prepare condition y (mask + latents) - T = clip_len_actual + 1 + T = segment_frame_length_actual + 1 # Encode condition y if mask_reft_len > 0: @@ -449,9 +455,9 @@ def __call__( y_reft = torch.cat([msk_reft, y_reft], dim=0) # Condition mask and latents - condition_msk_y = get_i2v_mask(lat_t_cond, latent_h, latent_w, clip_len_actual, device=device).to( - self.transformer.dtype - ) + condition_msk_y = get_i2v_mask( + lat_t_cond, latent_h, latent_w, segment_frame_length_actual, device=device + ).to(self.transformer.dtype) cond_lat_0 = condition_latents[0] if condition_latents.ndim == 5 else condition_latents condition_y = torch.cat([condition_msk_y, cond_lat_0], dim=0) @@ -489,7 +495,7 @@ def __call__( "seq_len": max_seq_len, "clip_fea": clip_context, "y": [y], - "origin_len": clip_len_actual, + "origin_len": segment_frame_length_actual, "origin_area": [actual_h, actual_w], } @@ -507,7 +513,7 @@ def __call__( "seq_len": max_seq_len, "clip_fea": clip_context, "y": [y], - "origin_len": clip_len_actual, + "origin_len": segment_frame_length_actual, "origin_area": [actual_h, actual_w], "is_uncondtion": True, } From e0291f88d94f59c4ed2a77cf5070f5d5ebb78283 Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Tue, 11 Aug 2026 23:16:37 +0000 Subject: [PATCH 09/19] Refine Wan-Animate-2 modular blocks: in-loop VAE encode, v1-style tensor names - Move the driving-video VAE encode back into the segment loop (per-segment causal encode, symmetric with the in-loop decode); the size check between the image/video preprocess outputs runs once in prepare_segments - Rename the research-code conditioning tensors to the merged Wan-Animate v1 names: y_ref -> reference_image_latents, y -> reference_latents, y_reft -> prev_segment_cond_latents; condition_latents/condition_y -> driving_video_latents/driving_video_condition - Rename the preprocessed video state to driving_video_pixels; derive latent/pixel dims from tensors instead of passing latent_height/latent_width - Collapse the four crop_* ints into a single crop_region tuple - Default height/width on the video preprocess step for standalone use; clarifying comments (zigzag padding, loop-carried state) Co-Authored-By: Claude Fable 5 --- .../wan_animate_2/before_denoise.py | 34 +++- .../wan_animate_2/decoders.py | 34 +--- .../wan_animate_2/denoise.py | 159 +++++++++++----- .../wan_animate_2/encoders.py | 169 ++++++------------ .../modular_blocks_wan_animate_2.py | 113 +++--------- .../modular_blocks_wan_animate_2_distilled.py | 34 ++-- 6 files changed, 239 insertions(+), 304 deletions(-) diff --git a/src/diffusers/modular_pipelines/wan_animate_2/before_denoise.py b/src/diffusers/modular_pipelines/wan_animate_2/before_denoise.py index 583899c87f8f..39c024a32dca 100644 --- a/src/diffusers/modular_pipelines/wan_animate_2/before_denoise.py +++ b/src/diffusers/modular_pipelines/wan_animate_2/before_denoise.py @@ -40,8 +40,19 @@ def description(self) -> str: def inputs(self) -> list[InputParam]: return [ InputParam("segment_frame_length", type_hint=int, default=81), - InputParam("latent_height", type_hint=int, required=True), - InputParam("latent_width", type_hint=int, required=True), + 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 @@ -69,16 +80,27 @@ def intermediate_outputs(self) -> list[OutputParam]: 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, block_state.latent_height, block_state.latent_width] + 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, block_state.latent_height // 2, block_state.latent_width // 2])) - ) + 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) diff --git a/src/diffusers/modular_pipelines/wan_animate_2/decoders.py b/src/diffusers/modular_pipelines/wan_animate_2/decoders.py index 7b52b0f2e886..5317ee85e67a 100644 --- a/src/diffusers/modular_pipelines/wan_animate_2/decoders.py +++ b/src/diffusers/modular_pipelines/wan_animate_2/decoders.py @@ -65,28 +65,11 @@ def inputs(self) -> list[InputParam]: description="Number of frames in the driving video before zigzag padding; the output is trimmed to it", ), InputParam( - "crop_top", + "crop_region", required=True, - type_hint=int, - description="Top edge of the reference image content inside the letterboxed frame", - ), - InputParam( - "crop_left", - required=True, - type_hint=int, - description="Left edge of the reference image content inside the letterboxed frame", - ), - InputParam( - "crop_height", - required=True, - type_hint=int, - description="Height of the reference image content inside the letterboxed frame", - ), - InputParam( - "crop_width", - required=True, - type_hint=int, - description="Width of the reference image content inside the letterboxed frame", + 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" @@ -108,13 +91,8 @@ 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] - video = video[ - :, - :, - :, - block_state.crop_top : block_state.crop_top + block_state.crop_height, - block_state.crop_left : block_state.crop_left + block_state.crop_width, - ] + 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) diff --git a/src/diffusers/modular_pipelines/wan_animate_2/denoise.py b/src/diffusers/modular_pipelines/wan_animate_2/denoise.py index aa845be8815a..4f42791f98f1 100644 --- a/src/diffusers/modular_pipelines/wan_animate_2/denoise.py +++ b/src/diffusers/modular_pipelines/wan_animate_2/denoise.py @@ -50,15 +50,86 @@ def decode_vae(vae: AutoencoderKLWan, latents: torch.Tensor) -> torch.Tensor: # ======================================== +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), + InputParam("effective_segment", required=True, type_hint=int), + InputParam("segment_frame_length", type_hint=int, default=81), + ] + + @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 `y`: the previous " + "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 `y_ref`. This is how motion continuity crosses segment boundaries — in pixel space, " + "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`." ) @@ -72,20 +143,16 @@ def expected_components(self) -> list[ComponentSpec]: @property def inputs(self) -> list[InputParam]: return [ - InputParam("y_ref", required=True, type_hint=torch.Tensor), + InputParam("reference_image_latents", required=True, type_hint=torch.Tensor), InputParam("segment_frame_length", type_hint=int, default=81), InputParam("prev_segment_conditioning_frames", type_hint=int, default=1), - InputParam("height", required=True, type_hint=int), - InputParam("width", required=True, type_hint=int), - InputParam("latent_height", required=True, type_hint=int), - InputParam("latent_width", required=True, type_hint=int), ] @property def intermediate_outputs(self) -> list[OutputParam]: return [ OutputParam( - "y", + "reference_latents", type_hint=torch.Tensor, description="The full conditioning tensor: reference half stacked over the segment half", ), @@ -97,30 +164,36 @@ def __call__(self, components, block_state: BlockState, k: int): # 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=(block_state.height, block_state.width), mode="bicubic" - ).permute(1, 0, 2, 3) + 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, block_state.height, block_state.width, device=device), + torch.zeros(3, num_frames - mask_len - 1, height, width, device=device), ], dim=1, ) else: - cond_pixels = torch.zeros(3, num_frames - 1, block_state.height, block_state.width, device=device) + cond_pixels = torch.zeros(3, num_frames - 1, height, width, device=device) - y_reft = encode_vae(components.vae, cond_pixels.unsqueeze(0)).squeeze(0) - mask_reft = get_i2v_mask( - y_reft.shape[1], block_state.latent_height, block_state.latent_width, mask_len, device=device - ).to(y_reft.dtype) - y_reft = torch.cat([mask_reft, y_reft], dim=0) + 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.y = torch.cat([block_state.y_ref, y_reft], dim=1) + block_state.reference_latents = torch.cat( + [block_state.reference_image_latents, prev_segment_cond_latents], dim=1 + ) return components, block_state @@ -146,9 +219,7 @@ def expected_components(self) -> list[ComponentSpec]: def inputs(self) -> list[InputParam]: return [ InputParam("generator"), - InputParam("y", required=True, type_hint=torch.Tensor), - InputParam("latent_height", required=True, type_hint=int), - InputParam("latent_width", required=True, type_hint=int), + InputParam("reference_latents", required=True, type_hint=torch.Tensor), ] @property @@ -168,9 +239,9 @@ def __call__(self, components, block_state: BlockState, k: int): block_state.latents = torch.randn( components.num_channels_latents, - block_state.y.shape[1], - block_state.latent_height, - block_state.latent_width, + block_state.reference_latents.shape[1], + block_state.reference_latents.shape[-2], + block_state.reference_latents.shape[-1], dtype=torch.float32, device=device, generator=block_state.generator, @@ -240,18 +311,8 @@ def expected_components(self) -> list[ComponentSpec]: @property def inputs(self) -> list[InputParam]: return [ - InputParam( - "condition_latents", - required=True, - type_hint=torch.Tensor, - description="VAE latents of every segment's driving-video slice, `[num_segments, 16, T, H', W']`", - ), - InputParam( - "condition_y", - required=True, - type_hint=torch.Tensor, - description="i2v mask + driving-slice latents per segment, `[num_segments, 20, T, H', W']`", - ), + InputParam("driving_video_latents", required=True, type_hint=torch.Tensor), + InputParam("driving_video_condition", required=True, type_hint=torch.Tensor), InputParam("condition_clip_context", required=True, type_hint=torch.Tensor), InputParam("prompt_ref_embeds", required=True, type_hint=torch.Tensor), InputParam("kv_cache", required=True, type_hint=WanAnimate2KVCache), @@ -267,11 +328,11 @@ def __call__(self, components, block_state: BlockState, k: int): t_ref = torch.tensor([block_state.timesteps[0].item()], device=device, dtype=transformer_dtype) components.transformer( - [block_state.condition_latents[k].to(transformer_dtype)], + [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.condition_y[k].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, @@ -315,7 +376,7 @@ def expected_components(self) -> list[ComponentSpec]: def inputs(self) -> list[InputParam]: return [ InputParam("latents", required=True, type_hint=torch.Tensor), - InputParam("y", required=True, type_hint=torch.Tensor), + InputParam("reference_latents", required=True, type_hint=torch.Tensor), InputParam("kv_cache", required=True, type_hint=WanAnimate2KVCache), InputParam("timesteps", required=True, type_hint=torch.Tensor), InputParam("num_inference_steps", type_hint=int, default=40), @@ -371,7 +432,7 @@ def __call__(self, components, block_state: BlockState, k: int): [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.y.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, @@ -498,8 +559,8 @@ class WanAnimate2SegmentLoopWrapper(LoopSequentialPipelineBlocks): def description(self) -> str: return ( "Pipeline block that iterates over the driving video's segments. At each segment it runs sub-blocks " - "for preparation, reference extraction, denoising, and decoding; each segment conditions on the " - "previous one's decoded tail frames." + "for per-segment encoding, preparation, reference extraction, denoising, and decoding; each segment " + "conditions on the previous one's decoded tail frames." ) @property @@ -522,6 +583,9 @@ def loop_intermediate_outputs(self) -> list[OutputParam]: 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 @@ -539,6 +603,7 @@ def __call__(self, components, state: PipelineState) -> PipelineState: class WanAnimate2DenoiseStep(WanAnimate2SegmentLoopWrapper): block_classes = [ + WanAnimate2SegmentVaeEncoderStep, WanAnimate2SegmentPrevFramesStep, WanAnimate2SegmentPrepareStep, WanAnimate2SegmentSchedulerResetStep, @@ -547,6 +612,7 @@ class WanAnimate2DenoiseStep(WanAnimate2SegmentLoopWrapper): WanAnimate2SegmentDecodeStep, ] block_names = [ + "vae_encoder", "prev_frames", "prepare", "scheduler_reset", @@ -559,12 +625,14 @@ class WanAnimate2DenoiseStep(WanAnimate2SegmentLoopWrapper): def description(self) -> str: return ( "Segment denoise step that iterates over the driving video's segments.\n" - "At each segment: prev_frames -> prepare -> scheduler_reset -> ref_extract -> denoise_inner -> decode." + "At each segment: vae_encoder -> prev_frames -> prepare -> scheduler_reset -> ref_extract -> " + "denoise_inner -> decode." ) class WanAnimate2DistilledDenoiseStep(WanAnimate2SegmentLoopWrapper): block_classes = [ + WanAnimate2SegmentVaeEncoderStep, WanAnimate2SegmentPrevFramesStep, WanAnimate2SegmentPrepareStep, WanAnimate2SegmentSchedulerResetStep, @@ -573,6 +641,7 @@ class WanAnimate2DistilledDenoiseStep(WanAnimate2SegmentLoopWrapper): WanAnimate2SegmentDecodeStep, ] block_names = [ + "vae_encoder", "prev_frames", "prepare", "scheduler_reset", @@ -585,6 +654,6 @@ class WanAnimate2DistilledDenoiseStep(WanAnimate2SegmentLoopWrapper): def description(self) -> str: return ( "Segment denoise step for the distilled model that iterates over the driving video's segments.\n" - "At each segment: prev_frames -> prepare -> scheduler_reset -> ref_extract -> " + "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 index 907f1de0bfc5..90cced247d18 100644 --- a/src/diffusers/modular_pipelines/wan_animate_2/encoders.py +++ b/src/diffusers/modular_pipelines/wan_animate_2/encoders.py @@ -264,10 +264,11 @@ def intermediate_outputs(self) -> list[OutputParam]: type_hint=torch.Tensor, description="The letterboxed reference image as a `[1, 3, H, W]` tensor in `[-1, 1]`", ), - OutputParam("crop_top", type_hint=int, description="Top edge of the content inside the letterbox"), - OutputParam("crop_left", type_hint=int, description="Left edge of the content inside the letterbox"), - OutputParam("crop_height", type_hint=int, description="Height of the content inside the letterbox"), - OutputParam("crop_width", type_hint=int, description="Width of the content inside the letterbox"), + 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() @@ -284,14 +285,11 @@ def __call__(self, components, state: PipelineState) -> PipelineState: block_state.width = int(math.sqrt(max_area / aspect_ratio)) // mod_value * mod_value height, width = block_state.height, block_state.width - block_state.crop_width = ( - width if width / height < image_width / image_height else image_width * height // image_height - ) - block_state.crop_height = ( - height if width / height >= image_width / image_height else image_height * width // image_width - ) - block_state.crop_top = (height - block_state.crop_height) // 2 - block_state.crop_left = (width - block_state.crop_width) // 2 + 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" @@ -332,7 +330,7 @@ def inputs(self) -> list[InputParam]: required=True, type_hint=list[PIL.Image.Image], description="The driving video that provides the motion, in any format accepted by " - "`VideoProcessor.preprocess_video`. Overwritten with the preprocessed `[1, 3, T, H, W]` tensor.", + "`VideoProcessor.preprocess_video`.", ), InputParam( "driving_video_fps", @@ -354,13 +352,30 @@ def inputs(self) -> list[InputParam]: default=1, description="The number of conditioning frames carried over from the previous segment", ), - InputParam("height", type_hint=int, required=True), - InputParam("width", type_hint=int, required=True), + 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, @@ -380,29 +395,36 @@ def __call__(self, components, state: PipelineState) -> PipelineState: 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: - last_segment_frames = (real_frame_len - block_state.prev_segment_conditioning_frames) % effective_segment + leftover_frames = (real_frame_len - block_state.prev_segment_conditioning_frames) % effective_segment else: - last_segment_frames = 0 - num_padding = effective_segment - last_segment_frames if last_segment_frames > 0 else 0 + 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 = driving_video + 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 = ( @@ -481,10 +503,10 @@ def expected_components(self) -> list[ComponentSpec]: def inputs(self) -> list[InputParam]: return [ InputParam( - "driving_video", + "driving_video_pixels", required=True, type_hint=torch.Tensor, - description="The preprocessed driving video `[1, 3, T, H, W]`", + description="The preprocessed driving video `[1, 3, T, H, W]`, from the video preprocess step", ), ] @@ -504,7 +526,7 @@ def __call__(self, components, state: PipelineState) -> PipelineState: device = components._execution_device block_state.condition_clip_context = clip_visual_encode( - components.image_encoder, block_state.driving_video[0, :, 0], device, components.image_encoder.dtype + components.image_encoder, block_state.driving_video_pixels[0, :, 0], device, components.image_encoder.dtype ) self.set_block_state(state, block_state) @@ -523,7 +545,7 @@ class WanAnimate2ImageVaeEncoderStep(ModularPipelineBlocks): 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 `y`" + "on top, producing the reference half of the conditioning tensor `reference_latents`" ) @property @@ -536,20 +558,16 @@ def expected_components(self) -> list[ComponentSpec]: def inputs(self) -> list[InputParam]: return [ InputParam("image_pixels", required=True, type_hint=torch.Tensor), - InputParam("height", type_hint=int, required=True), - InputParam("width", type_hint=int, required=True), ] @property def intermediate_outputs(self) -> list[OutputParam]: return [ OutputParam( - "y_ref", + "reference_image_latents", type_hint=torch.Tensor, description="i2v mask + reference image latents, `[20, 1, latent_height, latent_width]`", ), - OutputParam("latent_height", type_hint=int), - OutputParam("latent_width", type_hint=int), ] @torch.no_grad() @@ -560,97 +578,12 @@ def __call__(self, components, state: PipelineState) -> PipelineState: ref_latents = encode_vae(components.vae, block_state.image_pixels.unsqueeze(2)) - block_state.latent_height = block_state.height // components.vae_scale_factor_spatial - block_state.latent_width = block_state.width // components.vae_scale_factor_spatial - - mask_ref = get_i2v_mask(1, block_state.latent_height, block_state.latent_width, 1, device=device).to( - ref_latents.dtype - ) - block_state.y_ref = torch.cat([mask_ref, ref_latents[0]], dim=0) - - self.set_block_state(state, block_state) - return components, state - - -class WanAnimate2VideoVaeEncoderStep(ModularPipelineBlocks): - model_name = "wan-animate-2" - - @property - def description(self) -> str: - return ( - "VAE Encoder step that encodes every 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." - ) - - @property - def expected_components(self) -> list[ComponentSpec]: - return [ - ComponentSpec("vae", AutoencoderKLWan), - ] - - @property - def inputs(self) -> list[InputParam]: - return [ - InputParam( - "driving_video", - required=True, - type_hint=torch.Tensor, - description="The preprocessed driving video `[1, 3, T, H, W]`", - ), - InputParam("num_segments", required=True, type_hint=int), - InputParam("effective_segment", required=True, type_hint=int), - InputParam("segment_frame_length", type_hint=int, default=81), - InputParam("latent_height", required=True, type_hint=int), - InputParam("latent_width", required=True, type_hint=int), - ] + height, width = block_state.image_pixels.shape[-2:] + latent_height = height // components.vae_scale_factor_spatial + latent_width = width // components.vae_scale_factor_spatial - @property - def intermediate_outputs(self) -> list[OutputParam]: - return [ - OutputParam( - "condition_latents", - type_hint=torch.Tensor, - description="VAE latents of every segment's driving-video slice, `[num_segments, 16, T, H', W']`", - ), - OutputParam( - "condition_y", - type_hint=torch.Tensor, - description=( - "i2v mask + driving-slice latents per segment, `[num_segments, 20, T, H', W']`, conditioning " - "the reference-extraction pass" - ), - ), - ] - - @torch.no_grad() - def __call__(self, components, state: PipelineState) -> PipelineState: - block_state = self.get_block_state(state) - - device = components._execution_device - - condition_latents = [] - for k in range(block_state.num_segments): - start = k * block_state.effective_segment - condition_latents.append( - encode_vae( - components.vae, block_state.driving_video[:, :, start : start + block_state.segment_frame_length] - ) - ) - block_state.condition_latents = torch.cat(condition_latents, dim=0) - - # After zigzag padding every segment is exactly `segment_frame_length` frames, so one mask fits all. - condition_mask = get_i2v_mask( - block_state.condition_latents.shape[2], - block_state.latent_height, - block_state.latent_width, - block_state.segment_frame_length, - device=device, - ).to(block_state.condition_latents.dtype) - block_state.condition_y = torch.stack( - [torch.cat([condition_mask, latents], dim=0) for latents in block_state.condition_latents] - ) + 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 index b58a474acde1..a1b4d7e2b3a2 100644 --- 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 @@ -25,7 +25,6 @@ WanAnimate2ProcessVideosInputStep, WanAnimate2TextEncoderStep, WanAnimate2VideoClipEncoderStep, - WanAnimate2VideoVaeEncoderStep, ) @@ -66,14 +65,8 @@ class WanAnimate2ImageEncodeStep(SequentialPipelineBlocks): Outputs: image_pixels (`Tensor`): The letterboxed reference image as a `[1, 3, H, W]` tensor in `[-1, 1]` - crop_top (`int`): - Top edge of the content inside the letterbox - crop_left (`int`): - Left edge of the content inside the letterbox - crop_height (`int`): - Height of the content inside the letterbox - crop_width (`int`): - Width of the content inside the letterbox + 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 """ @@ -110,7 +103,6 @@ class WanAnimate2VideoEncodeStep(SequentialPipelineBlocks): Inputs: driving_video (`list`): The driving video that provides the motion, in any format accepted by `VideoProcessor.preprocess_video`. - Overwritten with the preprocessed `[1, 3, T, H, W]` tensor. 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. @@ -120,12 +112,15 @@ class WanAnimate2VideoEncodeStep(SequentialPipelineBlocks): 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`): - TODO: Add description. - width (`int`): - TODO: Add description. + 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`): @@ -148,65 +143,6 @@ def description(self): ) -WanAnimate2VaeEncoderBlocks = InsertableDict( - [ - ("image_encode", WanAnimate2ImageVaeEncoderStep()), - ("video_encode", WanAnimate2VideoVaeEncoderStep()), - ] -) - - -# auto_docstring -class WanAnimate2VaeEncodeStep(SequentialPipelineBlocks): - """ - VAE encoder step that encodes the letterboxed reference image into the reference half of the conditioning tensor `y_ref`, and every segment's slice of the driving video into the reference-extraction conditioning `condition_latents` / `condition_y`. - - Components: - vae (`AutoencoderKLWan`) - - Inputs: - image_pixels (`Tensor`): - TODO: Add description. - height (`int`): - TODO: Add description. - width (`int`): - TODO: Add description. - driving_video (`Tensor`): - The preprocessed driving video `[1, 3, T, H, W]` - num_segments (`int`): - TODO: Add description. - effective_segment (`int`): - TODO: Add description. - segment_frame_length (`int`, *optional*, defaults to 81): - TODO: Add description. - - Outputs: - y_ref (`Tensor`): - i2v mask + reference image latents, `[20, 1, latent_height, latent_width]` - latent_height (`int`): - TODO: Add description. - latent_width (`int`): - TODO: Add description. - condition_latents (`Tensor`): - VAE latents of every segment's driving-video slice, `[num_segments, 16, T, H', W']` - condition_y (`Tensor`): - i2v mask + driving-slice latents per segment, `[num_segments, 20, T, H', W']`, conditioning the - reference-extraction pass - """ - - model_name = "wan-animate-2" - block_classes = WanAnimate2VaeEncoderBlocks.values() - block_names = WanAnimate2VaeEncoderBlocks.keys() - - @property - def description(self): - return ( - "VAE encoder step that encodes the letterboxed reference image into the reference half of the " - "conditioning tensor `y_ref`, and every segment's slice of the driving video into the " - "reference-extraction conditioning `condition_latents` / `condition_y`." - ) - - # ==================== # 2. Core denoise # ==================== @@ -234,32 +170,28 @@ class WanAnimate2CoreDenoiseStep(SequentialPipelineBlocks): Inputs: segment_frame_length (`int`, *optional*, defaults to 81): TODO: Add description. - latent_height (`int`): - TODO: Add description. - latent_width (`int`): - TODO: Add description. + 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`): TODO: Add description. - y_ref (`Tensor`): + effective_segment (`int`): TODO: Add description. prev_segment_conditioning_frames (`int`, *optional*, defaults to 1): TODO: Add description. - height (`int`): - TODO: Add description. - width (`int`): - TODO: Add description. generator (`None`, *optional*): TODO: Add description. num_inference_steps (`int`, *optional*, defaults to 40): TODO: Add description. - condition_latents (`Tensor`): - VAE latents of every segment's driving-video slice, `[num_segments, 16, T, H', W']` - condition_y (`Tensor`): - i2v mask + driving-slice latents per segment, `[num_segments, 20, T, H', W']` condition_clip_context (`Tensor`): TODO: Add description. prompt_ref_embeds (`Tensor`): TODO: Add description. + height (`int`): + TODO: Add description. + width (`int`): + TODO: Add description. prompt_embeds (`Tensor`): text embeddings used to guide the image generation. Can be generated from text_encoder step. negative_prompt_embeds (`Tensor`, *optional*): @@ -275,7 +207,11 @@ class WanAnimate2CoreDenoiseStep(SequentialPipelineBlocks): Packed sequence length of the generation tokens max_seq_len_ref (`int`): Packed sequence length of the reference tokens - y (`Tensor`): + driving_video_latents (`Tensor`): + VAE latents of this segment's driving-video slice + driving_video_condition (`Tensor`): + i2v mask + driving-slice latents, conditioning the reference-extraction pass + reference_latents (`Tensor`): The full conditioning tensor: reference half stacked over the segment half latents (`Tensor`): This segment's initial noise @@ -312,7 +248,7 @@ def description(self): ("text_encoder", WanAnimate2TextEncoderStep()), ("image_encoder", WanAnimate2ImageEncodeStep()), ("video_encoder", WanAnimate2VideoEncodeStep()), - ("vae_encoder", WanAnimate2VaeEncodeStep()), + ("vae_encoder", WanAnimate2ImageVaeEncoderStep()), ("denoise", WanAnimate2CoreDenoiseStep()), ("decode", WanAnimate2DecodeStep()), ] @@ -353,7 +289,6 @@ class WanAnimate2Blocks(SequentialPipelineBlocks): 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`. - Overwritten with the preprocessed `[1, 3, T, H, W]` tensor. 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. 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 index 5394538f3eb8..36fd9cafb75d 100644 --- 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 @@ -18,10 +18,9 @@ from .before_denoise import WanAnimate2PrepareSegmentsStep from .decoders import WanAnimate2DecodeStep from .denoise import WanAnimate2DistilledDenoiseStep -from .encoders import WanAnimate2TextEncoderStep +from .encoders import WanAnimate2ImageVaeEncoderStep, WanAnimate2TextEncoderStep from .modular_blocks_wan_animate_2 import ( WanAnimate2ImageEncodeStep, - WanAnimate2VaeEncodeStep, WanAnimate2VideoEncodeStep, ) @@ -51,32 +50,28 @@ class WanAnimate2DistilledCoreDenoiseStep(SequentialPipelineBlocks): Inputs: segment_frame_length (`int`, *optional*, defaults to 81): TODO: Add description. - latent_height (`int`): - TODO: Add description. - latent_width (`int`): - TODO: Add description. + 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`): TODO: Add description. - y_ref (`Tensor`): + effective_segment (`int`): TODO: Add description. prev_segment_conditioning_frames (`int`, *optional*, defaults to 1): TODO: Add description. - height (`int`): - TODO: Add description. - width (`int`): - TODO: Add description. generator (`None`, *optional*): TODO: Add description. num_inference_steps (`int`, *optional*, defaults to 40): TODO: Add description. - condition_latents (`Tensor`): - VAE latents of every segment's driving-video slice, `[num_segments, 16, T, H', W']` - condition_y (`Tensor`): - i2v mask + driving-slice latents per segment, `[num_segments, 20, T, H', W']` condition_clip_context (`Tensor`): TODO: Add description. prompt_ref_embeds (`Tensor`): TODO: Add description. + height (`int`): + TODO: Add description. + width (`int`): + TODO: Add description. prompt_embeds (`Tensor`): text embeddings used to guide the image generation. Can be generated from text_encoder step. negative_prompt_embeds (`Tensor`, *optional*): @@ -92,7 +87,11 @@ class WanAnimate2DistilledCoreDenoiseStep(SequentialPipelineBlocks): Packed sequence length of the generation tokens max_seq_len_ref (`int`): Packed sequence length of the reference tokens - y (`Tensor`): + driving_video_latents (`Tensor`): + VAE latents of this segment's driving-video slice + driving_video_condition (`Tensor`): + i2v mask + driving-slice latents, conditioning the reference-extraction pass + reference_latents (`Tensor`): The full conditioning tensor: reference half stacked over the segment half latents (`Tensor`): This segment's initial noise @@ -123,7 +122,7 @@ def description(self): ("text_encoder", WanAnimate2TextEncoderStep()), ("image_encoder", WanAnimate2ImageEncodeStep()), ("video_encoder", WanAnimate2VideoEncodeStep()), - ("vae_encoder", WanAnimate2VaeEncodeStep()), + ("vae_encoder", WanAnimate2ImageVaeEncoderStep()), ("denoise", WanAnimate2DistilledCoreDenoiseStep()), ("decode", WanAnimate2DecodeStep()), ] @@ -164,7 +163,6 @@ class WanAnimate2DistilledBlocks(SequentialPipelineBlocks): 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`. - Overwritten with the preprocessed `[1, 3, T, H, W]` tensor. 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. From 1b81e877a8132a8f896fd334af6769a55b6bbc85 Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Wed, 12 Aug 2026 02:27:07 +0000 Subject: [PATCH 10/19] Add Wan-Animate-2 modular tests, fill docstrings, self-assemble blocksets MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - tests/modular_pipelines/wan_animate_2/: ModularPipelineTesterMixin + ModularGuiderTesterMixin against YiYiXu/tiny-wan-animate-2-modular and -distilled-modular (25 passed / 15 skipped); batch tests skipped (the pipeline is unbatched), guider test threshold lowered with rationale - Use randn_tensor for segment noise in both pipelines so CPU generators work; CUDA-generator path unchanged (parity re-verified bit-identical) - Describe every InputParam (templates where available); no auto-docstring TODOs remain - Narrow the core denoise steps' outputs to segment_frames — the only product the decode step consumes - Make the distilled blockset file self-contained: it assembles its own image/video encoder groups from the leaf blocks instead of importing the base blockset's classes Co-Authored-By: Claude Fable 5 --- .../wan_animate_2/before_denoise.py | 7 +- .../wan_animate_2/denoise.py | 228 +++++++++++++++--- .../wan_animate_2/encoders.py | 22 +- .../modular_blocks_wan_animate_2.py | 82 +++---- .../modular_blocks_wan_animate_2_distilled.py | 210 ++++++++++++---- .../pipelines/wan/pipeline_wan_animate_2.py | 12 +- .../wan_animate_2/__init__.py | 0 .../test_modular_pipeline_wan_animate_2.py | 85 +++++++ 8 files changed, 500 insertions(+), 146 deletions(-) create mode 100644 tests/modular_pipelines/wan_animate_2/__init__.py create mode 100644 tests/modular_pipelines/wan_animate_2/test_modular_pipeline_wan_animate_2.py diff --git a/src/diffusers/modular_pipelines/wan_animate_2/before_denoise.py b/src/diffusers/modular_pipelines/wan_animate_2/before_denoise.py index 39c024a32dca..0ad038e8ccaa 100644 --- a/src/diffusers/modular_pipelines/wan_animate_2/before_denoise.py +++ b/src/diffusers/modular_pipelines/wan_animate_2/before_denoise.py @@ -39,7 +39,12 @@ def description(self) -> str: @property def inputs(self) -> list[InputParam]: return [ - InputParam("segment_frame_length", type_hint=int, default=81), + InputParam( + "segment_frame_length", + type_hint=int, + default=81, + description="The number of frames in each inference segment", + ), InputParam( "reference_image_latents", required=True, diff --git a/src/diffusers/modular_pipelines/wan_animate_2/denoise.py b/src/diffusers/modular_pipelines/wan_animate_2/denoise.py index 4f42791f98f1..7a97c1c13db5 100644 --- a/src/diffusers/modular_pipelines/wan_animate_2/denoise.py +++ b/src/diffusers/modular_pipelines/wan_animate_2/denoise.py @@ -24,6 +24,7 @@ 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 @@ -78,9 +79,24 @@ def inputs(self) -> list[InputParam]: 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), - InputParam("effective_segment", required=True, type_hint=int), - InputParam("segment_frame_length", type_hint=int, default=81), + 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 @@ -143,9 +159,24 @@ def expected_components(self) -> list[ComponentSpec]: @property def inputs(self) -> list[InputParam]: return [ - InputParam("reference_image_latents", required=True, type_hint=torch.Tensor), - InputParam("segment_frame_length", type_hint=int, default=81), - InputParam("prev_segment_conditioning_frames", type_hint=int, default=1), + 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 @@ -218,8 +249,13 @@ def expected_components(self) -> list[ComponentSpec]: @property def inputs(self) -> list[InputParam]: return [ - InputParam("generator"), - InputParam("reference_latents", required=True, type_hint=torch.Tensor), + 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 @@ -237,14 +273,16 @@ def intermediate_outputs(self) -> list[OutputParam]: def __call__(self, components, block_state: BlockState, k: int): device = components._execution_device - block_state.latents = torch.randn( - components.num_channels_latents, - block_state.reference_latents.shape[1], - block_state.reference_latents.shape[-2], - block_state.reference_latents.shape[-1], - dtype=torch.float32, - device=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) @@ -271,7 +309,7 @@ def expected_components(self) -> list[ComponentSpec]: @property def inputs(self) -> list[InputParam]: return [ - InputParam("num_inference_steps", type_hint=int, default=40), + InputParam.template("num_inference_steps", default=40), ] @property @@ -311,14 +349,54 @@ def expected_components(self) -> list[ComponentSpec]: @property def inputs(self) -> list[InputParam]: return [ - InputParam("driving_video_latents", required=True, type_hint=torch.Tensor), - InputParam("driving_video_condition", required=True, type_hint=torch.Tensor), - InputParam("condition_clip_context", required=True, type_hint=torch.Tensor), - InputParam("prompt_ref_embeds", required=True, type_hint=torch.Tensor), - InputParam("kv_cache", required=True, type_hint=WanAnimate2KVCache), - InputParam("timesteps", required=True, type_hint=torch.Tensor), - InputParam("max_seq_len_ref", required=True, type_hint=int), - InputParam("grid_sizes_ref", required=True, type_hint=torch.Tensor), + 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() @@ -375,18 +453,68 @@ def expected_components(self) -> list[ComponentSpec]: @property def inputs(self) -> list[InputParam]: return [ - InputParam("latents", required=True, type_hint=torch.Tensor), - InputParam("reference_latents", required=True, type_hint=torch.Tensor), - InputParam("kv_cache", required=True, type_hint=WanAnimate2KVCache), - InputParam("timesteps", required=True, type_hint=torch.Tensor), - InputParam("num_inference_steps", type_hint=int, default=40), - InputParam("num_segments", required=True, type_hint=int), - InputParam("max_seq_len", required=True, type_hint=int), - InputParam("grid_sizes_ref", required=True, type_hint=torch.Tensor), - InputParam("segment_frame_length", type_hint=int, default=81), - InputParam("height", required=True, type_hint=int), - InputParam("width", required=True, type_hint=int), - InputParam("generator"), + 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"), @@ -512,9 +640,24 @@ def expected_components(self) -> list[ComponentSpec]: @property def inputs(self) -> list[InputParam]: return [ - InputParam("latents", required=True, type_hint=torch.Tensor), - InputParam("kv_cache", required=True, type_hint=WanAnimate2KVCache), - InputParam("prev_segment_conditioning_frames", type_hint=int, default=1), + 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 @@ -566,7 +709,12 @@ def description(self) -> str: @property def loop_inputs(self) -> list[InputParam]: return [ - InputParam("num_segments", required=True, type_hint=int), + InputParam( + "num_segments", + required=True, + type_hint=int, + description="Total number of segments in the driving video, from the video preprocess step", + ), ] @property diff --git a/src/diffusers/modular_pipelines/wan_animate_2/encoders.py b/src/diffusers/modular_pipelines/wan_animate_2/encoders.py index 90cced247d18..4100c9aa0181 100644 --- a/src/diffusers/modular_pipelines/wan_animate_2/encoders.py +++ b/src/diffusers/modular_pipelines/wan_animate_2/encoders.py @@ -143,15 +143,15 @@ def expected_components(self) -> list[ComponentSpec]: @property def inputs(self) -> list[InputParam]: return [ - InputParam("prompt", required=True, type_hint=str), - InputParam("negative_prompt", type_hint=str), + InputParam.template("prompt"), + InputParam.template("negative_prompt"), InputParam( "prompt_ref", default="人物动作的参考视频", type_hint=str, description="The reference prompt for the driving video context", ), - InputParam("max_sequence_length", default=512), + InputParam.template("max_sequence_length"), ] @property @@ -240,7 +240,7 @@ def expected_components(self) -> list[ComponentSpec]: @property def inputs(self) -> list[InputParam]: return [ - InputParam("image", type_hint=PIL.Image.Image, required=True), + InputParam.template("image", description="The reference image holding the character to animate."), InputParam( "height", type_hint=int, @@ -456,7 +456,12 @@ def expected_components(self) -> list[ComponentSpec]: @property def inputs(self) -> list[InputParam]: return [ - InputParam("image_pixels", required=True, type_hint=torch.Tensor), + 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 @@ -557,7 +562,12 @@ def expected_components(self) -> list[ComponentSpec]: @property def inputs(self) -> list[InputParam]: return [ - InputParam("image_pixels", required=True, type_hint=torch.Tensor), + 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 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 index a1b4d7e2b3a2..998e04f4f264 100644 --- 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 @@ -12,6 +12,8 @@ # 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 @@ -54,8 +56,8 @@ class WanAnimate2ImageEncodeStep(SequentialPipelineBlocks): image_encoder (`CLIPVisionModel`) Inputs: - image (`Image`): - TODO: Add description. + 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. @@ -169,29 +171,30 @@ class WanAnimate2CoreDenoiseStep(SequentialPipelineBlocks): Inputs: segment_frame_length (`int`, *optional*, defaults to 81): - TODO: Add description. + 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`): - TODO: Add description. + Total number of segments in the driving video, from the video preprocess step effective_segment (`int`): - TODO: Add description. + 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): - TODO: Add description. - generator (`None`, *optional*): - TODO: Add description. + 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): - TODO: Add description. + The number of denoising steps. condition_clip_context (`Tensor`): - TODO: Add description. + CLIP vision features of the driving video's first frame prompt_ref_embeds (`Tensor`): - TODO: Add description. + Text embeddings of the reference prompt, guiding the reference-extraction pass height (`int`): - TODO: Add description. + The resolved frame height in pixels width (`int`): - TODO: Add description. + 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*): @@ -200,29 +203,9 @@ class WanAnimate2CoreDenoiseStep(SequentialPipelineBlocks): conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. Outputs: - grid_sizes_ref (`Tensor`): - 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 - max_seq_len (`int`): - Packed sequence length of the generation tokens - max_seq_len_ref (`int`): - Packed sequence length of the reference tokens - driving_video_latents (`Tensor`): - VAE latents of this segment's driving-video slice - driving_video_condition (`Tensor`): - i2v mask + driving-slice latents, conditioning the reference-extraction pass - reference_latents (`Tensor`): - The full conditioning tensor: reference half stacked over the segment half - latents (`Tensor`): - This segment's initial noise - kv_cache (`WanAnimate2KVCache`): - Fresh per-segment cache for the reference K/V - timesteps (`Tensor`): - This segment's denoising timesteps - out_frames (`Tensor`): - This segment's decoded frames on device; the next segment conditions on its tail segment_frames (`list`): - Per-segment decoded frames on CPU, each `[1, 3, T, H, W]` + 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" @@ -237,6 +220,17 @@ def description(self): "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 @@ -273,15 +267,15 @@ class WanAnimate2Blocks(SequentialPipelineBlocks): Inputs: prompt (`str`): - TODO: Add description. + The prompt or prompts to guide image generation. negative_prompt (`str`, *optional*): - TODO: Add description. + 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 (`None`, *optional*, defaults to 512): - TODO: Add description. - image (`Image`): - TODO: Add description. + 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. @@ -298,10 +292,10 @@ class WanAnimate2Blocks(SequentialPipelineBlocks): 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 (`None`, *optional*): - TODO: Add description. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. num_inference_steps (`int`, *optional*, defaults to 40): - TODO: Add description. + 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): 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 index 36fd9cafb75d..a2a9869d27f7 100644 --- 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 @@ -12,22 +12,144 @@ # 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 WanAnimate2DistilledDenoiseStep -from .encoders import WanAnimate2ImageVaeEncoderStep, WanAnimate2TextEncoderStep -from .modular_blocks_wan_animate_2 import ( - WanAnimate2ImageEncodeStep, - WanAnimate2VideoEncodeStep, +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" + 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" + 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()), @@ -49,29 +171,30 @@ class WanAnimate2DistilledCoreDenoiseStep(SequentialPipelineBlocks): Inputs: segment_frame_length (`int`, *optional*, defaults to 81): - TODO: Add description. + 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`): - TODO: Add description. + Total number of segments in the driving video, from the video preprocess step effective_segment (`int`): - TODO: Add description. + 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): - TODO: Add description. - generator (`None`, *optional*): - TODO: Add description. + 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): - TODO: Add description. + The number of denoising steps. condition_clip_context (`Tensor`): - TODO: Add description. + CLIP vision features of the driving video's first frame prompt_ref_embeds (`Tensor`): - TODO: Add description. + Text embeddings of the reference prompt, guiding the reference-extraction pass height (`int`): - TODO: Add description. + The resolved frame height in pixels width (`int`): - TODO: Add description. + 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*): @@ -80,29 +203,9 @@ class WanAnimate2DistilledCoreDenoiseStep(SequentialPipelineBlocks): conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. Outputs: - grid_sizes_ref (`Tensor`): - 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 - max_seq_len (`int`): - Packed sequence length of the generation tokens - max_seq_len_ref (`int`): - Packed sequence length of the reference tokens - driving_video_latents (`Tensor`): - VAE latents of this segment's driving-video slice - driving_video_condition (`Tensor`): - i2v mask + driving-slice latents, conditioning the reference-extraction pass - reference_latents (`Tensor`): - The full conditioning tensor: reference half stacked over the segment half - latents (`Tensor`): - This segment's initial noise - kv_cache (`WanAnimate2KVCache`): - Fresh per-segment cache for the reference K/V - timesteps (`Tensor`): - This segment's denoising timesteps - out_frames (`Tensor`): - This segment's decoded frames on device; the next segment conditions on its tail segment_frames (`list`): - Per-segment decoded frames on CPU, each `[1, 3, T, H, W]` + 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" @@ -116,12 +219,23 @@ def description(self): "and runs the segment-by-segment denoising loop in few steps without classifier-free guidance." ) + @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", WanAnimate2ImageEncodeStep()), - ("video_encoder", WanAnimate2VideoEncodeStep()), + ("image_encoder", WanAnimate2DistilledImageEncodeStep()), + ("video_encoder", WanAnimate2DistilledVideoEncodeStep()), ("vae_encoder", WanAnimate2ImageVaeEncoderStep()), ("denoise", WanAnimate2DistilledCoreDenoiseStep()), ("decode", WanAnimate2DecodeStep()), @@ -147,15 +261,15 @@ class WanAnimate2DistilledBlocks(SequentialPipelineBlocks): Inputs: prompt (`str`): - TODO: Add description. + The prompt or prompts to guide image generation. negative_prompt (`str`, *optional*): - TODO: Add description. + 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 (`None`, *optional*, defaults to 512): - TODO: Add description. - image (`Image`): - TODO: Add description. + 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. @@ -172,10 +286,10 @@ class WanAnimate2DistilledBlocks(SequentialPipelineBlocks): 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 (`None`, *optional*): - TODO: Add description. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. num_inference_steps (`int`, *optional*, defaults to 40): - TODO: Add description. + 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): diff --git a/src/diffusers/pipelines/wan/pipeline_wan_animate_2.py b/src/diffusers/pipelines/wan/pipeline_wan_animate_2.py index 888b4191417d..91388ff091d5 100644 --- a/src/diffusers/pipelines/wan/pipeline_wan_animate_2.py +++ b/src/diffusers/pipelines/wan/pipeline_wan_animate_2.py @@ -27,6 +27,7 @@ from ...modular_pipelines.wan_animate_2.video_processor import WanAnimate2VideoProcessor from ...schedulers import SchedulerMixin from ...utils import logging +from ...utils.torch_utils import randn_tensor from ..pipeline_utils import DiffusionPipeline from .pipeline_output import WanPipelineOutput @@ -474,14 +475,11 @@ def __call__( # Noise latents temporal dim = y_ref(1) + y_reft/condition_y(T) = total y temporal dim lat_t_noise = y.shape[1] if y.ndim == 4 else y.shape[2] - noise = torch.randn( - 16, - lat_t_noise, - latent_h, - latent_w, - dtype=torch.float32, - device=device, + noise = randn_tensor( + (16, lat_t_noise, latent_h, latent_w), generator=generator, + device=device, + dtype=torch.float32, ) latents = [noise] 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..468c70accdd8 --- /dev/null +++ b/tests/modular_pipelines/wan_animate_2/test_modular_pipeline_wan_animate_2.py @@ -0,0 +1,85 @@ +# 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 + + +class TestWanAnimate2ModularPipelineFast(ModularPipelineTesterMixin, ModularGuiderTesterMixin): + pipeline_class = WanAnimate2ModularPipeline + pipeline_blocks_class = WanAnimate2Blocks + pretrained_model_name_or_path = "YiYiXu/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" + + 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 = "YiYiXu/tiny-wan-animate-2-distilled-modular" + + @pytest.mark.skip(reason="The distilled preset pins its guider to guidance_scale=1.0") + def test_guider_cfg(self): + pass From b9a15862ff4dae31e3dabd91908c55848be2eb54 Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Wed, 12 Aug 2026 02:37:14 +0000 Subject: [PATCH 11/19] Add Wan-Animate-2 docs (modular pipeline + transformer) Modular-only pipeline page (ModularPipeline.from_pretrained example with offloading + compile, both presets, area-based height/width semantics) and the transformer model page; toctree entries sorted. References the official Wan-AI hub ids, which will need the converted weights and modular_model_index.json before the examples run as written. Co-Authored-By: Claude Fable 5 --- docs/source/en/_toctree.yml | 4 + .../models/wan_animate_2_transformer_3d.md | 30 ++++++++ docs/source/en/api/pipelines/wan_animate_2.md | 77 +++++++++++++++++++ 3 files changed, 111 insertions(+) create mode 100644 docs/source/en/api/models/wan_animate_2_transformer_3d.md create mode 100644 docs/source/en/api/pipelines/wan_animate_2.md diff --git a/docs/source/en/_toctree.yml b/docs/source/en/_toctree.yml index 2d0bf5707ad9..a19a9ae02385 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 @@ -709,6 +711,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..e3b2ca3fcff3 --- /dev/null +++ b/docs/source/en/api/pipelines/wan_animate_2.md @@ -0,0 +1,77 @@ + + +# 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", + num_inference_steps=40, + 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 and pass `num_inference_steps=10`; no `guidance_scale` argument exists anywhere — guidance is owned by the pipeline's guider component. + +`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 From 1be4ed485b88195342df23df460dce850371352b Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Wed, 12 Aug 2026 02:59:55 +0000 Subject: [PATCH 12/19] Run make style / make quality Import sorting in the two __init__s and doc-builder docstring reflow; no behavior changes. make quality now exits clean. Co-Authored-By: Claude Fable 5 --- src/diffusers/__init__.py | 10 +-- src/diffusers/models/__init__.py | 2 +- .../transformers/transformer_wan_animate_2.py | 30 ++++----- .../modular_blocks_wan_animate_2.py | 64 +++++++++---------- .../modular_blocks_wan_animate_2_distilled.py | 64 +++++++++---------- .../wan_animate_2/video_processor.py | 12 ++-- .../pipelines/wan/pipeline_wan_animate_2.py | 27 ++++---- src/diffusers/utils/loading_utils.py | 4 +- 8 files changed, 100 insertions(+), 113 deletions(-) diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py index 36047ec5ca62..fa2f5ffe62f4 100644 --- a/src/diffusers/__init__.py +++ b/src/diffusers/__init__.py @@ -351,8 +351,8 @@ "UNetSpatioTemporalConditionModel", "UVit2DModel", "VQModel", - "WanAnimateTransformer3DModel", "WanAnimate2Transformer3DModel", + "WanAnimateTransformer3DModel", "WanTransformer3DModel", "WanVACETransformer3DModel", "ZImageControlNetModel", @@ -554,13 +554,13 @@ "Wan22Image2VideoBlocks", "Wan22Image2VideoModularPipeline", "Wan22ModularPipeline", - "WanBlocks", - "WanImage2VideoAutoBlocks", - "WanImage2VideoModularPipeline", "WanAnimate2Blocks", "WanAnimate2DistilledBlocks", "WanAnimate2DistilledModularPipeline", "WanAnimate2ModularPipeline", + "WanBlocks", + "WanImage2VideoAutoBlocks", + "WanImage2VideoModularPipeline", "WanModularPipeline", "ZImageAutoBlocks", "ZImageModularPipeline", @@ -848,8 +848,8 @@ "VisualClozeGenerationPipeline", "VisualClozePipeline", "VQDiffusionPipeline", - "WanAnimatePipeline", "WanAnimate2Pipeline", + "WanAnimatePipeline", "WanImageToVideoPipeline", "WanPipeline", "WanVACEPipeline", diff --git a/src/diffusers/models/__init__.py b/src/diffusers/models/__init__.py index 92dc0f15cc74..f5e0035c7517 100755 --- a/src/diffusers/models/__init__.py +++ b/src/diffusers/models/__init__.py @@ -284,8 +284,8 @@ T5FilmDecoder, Transformer2DModel, TransformerTemporalModel, - WanAnimateTransformer3DModel, WanAnimate2Transformer3DModel, + WanAnimateTransformer3DModel, WanTransformer3DModel, WanVACETransformer3DModel, ZImageTransformer2DModel, diff --git a/src/diffusers/models/transformers/transformer_wan_animate_2.py b/src/diffusers/models/transformers/transformer_wan_animate_2.py index 36309a0001fc..eef5cd168351 100644 --- a/src/diffusers/models/transformers/transformer_wan_animate_2.py +++ b/src/diffusers/models/transformers/transformer_wan_animate_2.py @@ -115,9 +115,8 @@ def _get_added_kv_projections(attn, encoder_hidden_states_img: torch.Tensor): 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)`. + 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): @@ -517,7 +516,9 @@ def forward( ) # 3. Feed-forward - norm_hidden_states = (self.norm2(hidden_states.float()) * (1 + c_scale_msa) + c_shift_msa).type_as(hidden_states) + 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) @@ -569,10 +570,9 @@ class WanAnimate2Transformer3DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, F 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. + (``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)`): @@ -787,23 +787,23 @@ def forward( 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. + `"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. + `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*), origin_area (`list[int]`, *optional*): - Frame count and spatial size of the full video, which the in-context block mask is built over. - Required under `kv_cache_mode="cached"`. + Frame count and spatial size 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`): 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 index 998e04f4f264..f77eb378c15b 100644 --- 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 @@ -49,18 +49,18 @@ # 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`. + 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`) + 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. + 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. @@ -96,18 +96,18 @@ def description(self): # 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`. + 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`) + 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. + 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): @@ -115,8 +115,8 @@ class WanAnimate2VideoEncodeStep(SequentialPipelineBlocks): 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. + 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`. @@ -161,13 +161,12 @@ def description(self): # 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. + 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`) + vae (`AutoencoderKLWan`) transformer (`WanAnimate2Transformer3DModel`) scheduler (`SchedulerMixin`) guider + (`ClassifierFreeGuidance`) Inputs: segment_frame_length (`int`, *optional*, defaults to 81): @@ -179,8 +178,8 @@ class WanAnimate2CoreDenoiseStep(SequentialPipelineBlocks): 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 + 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*): @@ -204,8 +203,8 @@ class WanAnimate2CoreDenoiseStep(SequentialPipelineBlocks): 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 + 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" @@ -252,18 +251,13 @@ def outputs(self): # 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. + 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`) + 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`): @@ -277,15 +271,15 @@ class WanAnimate2Blocks(SequentialPipelineBlocks): 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. + 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. + 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): 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 index a2a9869d27f7..660037342417 100644 --- 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 @@ -49,18 +49,18 @@ # 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`. + 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`) + 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. + 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. @@ -96,18 +96,18 @@ def description(self): # 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`. + 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`) + 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. + 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): @@ -115,8 +115,8 @@ class WanAnimate2DistilledVideoEncodeStep(SequentialPipelineBlocks): 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. + 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`. @@ -161,13 +161,12 @@ def description(self): # 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. + 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`) + vae (`AutoencoderKLWan`) transformer (`WanAnimate2Transformer3DModel`) scheduler (`SchedulerMixin`) guider + (`ClassifierFreeGuidance`) Inputs: segment_frame_length (`int`, *optional*, defaults to 81): @@ -179,8 +178,8 @@ class WanAnimate2DistilledCoreDenoiseStep(SequentialPipelineBlocks): 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 + 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*): @@ -204,8 +203,8 @@ class WanAnimate2DistilledCoreDenoiseStep(SequentialPipelineBlocks): 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 + 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" @@ -246,18 +245,13 @@ def outputs(self): # auto_docstring class WanAnimate2DistilledBlocks(SequentialPipelineBlocks): """ - Modular pipeline blocks for distilled Wan-Animate-2 character animation, sampling in few steps without classifier-free guidance. + 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`) + 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`): @@ -271,15 +265,15 @@ class WanAnimate2DistilledBlocks(SequentialPipelineBlocks): 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. + 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. + 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): diff --git a/src/diffusers/modular_pipelines/wan_animate_2/video_processor.py b/src/diffusers/modular_pipelines/wan_animate_2/video_processor.py index 742a55147063..18e041b35db7 100644 --- a/src/diffusers/modular_pipelines/wan_animate_2/video_processor.py +++ b/src/diffusers/modular_pipelines/wan_animate_2/video_processor.py @@ -23,12 +23,12 @@ 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. + 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 diff --git a/src/diffusers/pipelines/wan/pipeline_wan_animate_2.py b/src/diffusers/pipelines/wan/pipeline_wan_animate_2.py index 91388ff091d5..f9b01c3dc63f 100644 --- a/src/diffusers/pipelines/wan/pipeline_wan_animate_2.py +++ b/src/diffusers/pipelines/wan/pipeline_wan_animate_2.py @@ -77,10 +77,10 @@ class WanAnimate2Pipeline(DiffusionPipeline, WanLoraLoaderMixin): r""" Pipeline for character animation using Wan-Animate-2. - This pipeline takes a reference character image and a driving video, and generates a video where the character - is animated following the motion in the driving video. The model uses an in-context attention mechanism with - KV cache: a reference video is first encoded to cache K/V tensors, then the generation forward uses the cached - K/V with a block mask for frame-level sparse in-context attention. + This pipeline takes a reference character image and a driving video, and generates a video where the character is + animated following the motion in the driving video. The model uses an in-context attention mechanism with KV cache: + a reference video is first encoded to cache K/V tensors, then the generation forward uses the cached K/V with a + block mask for frame-level sparse in-context attention. Args: tokenizer ([`AutoTokenizer`]): @@ -92,9 +92,9 @@ class WanAnimate2Pipeline(DiffusionPipeline, WanLoraLoaderMixin): transformer ([`WanAnimate2Transformer3DModel`]): The Wan-Animate-2 transformer model. scheduler ([`SchedulerMixin`]): - A flow-matching scheduler to be used in combination with `transformer` to denoise the encoded latents. - The reference implementation samples with `DPMSolverMultistepScheduler` (`flow_shift=5.0`) for the base - model and `FlowMatchEulerDiscreteScheduler` (`shift=5.0`) for the distilled one. + A flow-matching scheduler to be used in combination with `transformer` to denoise the encoded latents. The + reference implementation samples with `DPMSolverMultistepScheduler` (`flow_shift=5.0`) for the base model + and `FlowMatchEulerDiscreteScheduler` (`shift=5.0`) for the distilled one. vae ([`AutoencoderKLWan`]): The Wan VAE model. """ @@ -262,8 +262,7 @@ def __call__( The reference character image. driving_video (`list[PIL.Image.Image]`, `np.ndarray` or `torch.Tensor`): The driving video that provides the motion, in any format accepted by - [`~video_processor.VideoProcessor.preprocess_video`]. Load one from disk with - [`~utils.load_video`]. + [`~video_processor.VideoProcessor.preprocess_video`]. Load one from disk with [`~utils.load_video`]. prompt (`str` or `list[str]`): The text prompt describing the character appearance and background. negative_prompt (`str` or `list[str]`, *optional*): @@ -271,9 +270,9 @@ def __call__( prompt_ref (`str`, defaults to `"人物动作的参考视频"`): The reference prompt for the driving video context. height (`int`, defaults to `800`): - Together with `width`, the target *area* (`height * width`) of the generated video. The aspect ratio - is taken from `image`, so the video is rarely exactly `height` x `width` — both dimensions are - rescaled to hit that area and then floored to a multiple of 16. + Together with `width`, the target *area* (`height * width`) of the generated video. The aspect ratio is + taken from `image`, so the video is rarely exactly `height` x `width` — both dimensions are rescaled to + hit that area and then floored to a multiple of 16. width (`int`, defaults to `640`): See `height`. segment_frame_length (`int`, defaults to `81`): @@ -281,8 +280,8 @@ def __call__( prev_segment_conditioning_frames (`int`, defaults to `1`): The number of conditioning frames from the previous segment. fps (`int`, defaults to `24`): - The frame rate the model generates at. `driving_video` is resampled to it when - `driving_video_fps` is given. + The frame rate the model generates at. `driving_video` is resampled to it when `driving_video_fps` is + given. driving_video_fps (`float`, *optional*): The frame rate `driving_video` was captured at — a list of frames does not carry it, so [`~utils.load_video`] will report it with `return_fps=True`. When set, the driving frames are diff --git a/src/diffusers/utils/loading_utils.py b/src/diffusers/utils/loading_utils.py index 7acfd6f6597f..ba528c7f857c 100644 --- a/src/diffusers/utils/loading_utils.py +++ b/src/diffusers/utils/loading_utils.py @@ -69,8 +69,8 @@ def load_video( 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. + 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]` or `tuple[list[PIL.Image.Image], float]`: From c3076098abdff3a436735c2c9858ccdad29c90b9 Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Wed, 12 Aug 2026 04:15:03 +0000 Subject: [PATCH 13/19] Point Wan-Animate-2 tests at the hf-internal-testing tiny repos Co-Authored-By: Claude Fable 5 --- .../wan_animate_2/test_modular_pipeline_wan_animate_2.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) 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 index 468c70accdd8..4b6048fdcb46 100644 --- 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 @@ -30,7 +30,7 @@ class TestWanAnimate2ModularPipelineFast(ModularPipelineTesterMixin, ModularGuiderTesterMixin): pipeline_class = WanAnimate2ModularPipeline pipeline_blocks_class = WanAnimate2Blocks - pretrained_model_name_or_path = "YiYiXu/tiny-wan-animate-2-modular" + pretrained_model_name_or_path = "hf-internal-testing/tiny-wan-animate-2-modular" params = frozenset(["prompt", "image", "driving_video"]) batch_params = frozenset() @@ -78,7 +78,7 @@ def test_guider_cfg(self): class TestWanAnimate2DistilledModularPipelineFast(TestWanAnimate2ModularPipelineFast): pipeline_class = WanAnimate2DistilledModularPipeline pipeline_blocks_class = WanAnimate2DistilledBlocks - pretrained_model_name_or_path = "YiYiXu/tiny-wan-animate-2-distilled-modular" + pretrained_model_name_or_path = "hf-internal-testing/tiny-wan-animate-2-distilled-modular" @pytest.mark.skip(reason="The distilled preset pins its guider to guidance_scale=1.0") def test_guider_cfg(self): From 6dc53de341169bbfae63502c180b2f788921dba2 Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Wed, 12 Aug 2026 04:38:46 +0000 Subject: [PATCH 14/19] Import the video processor lazily to break a circular import Under eager initialization (DIFFUSERS_SLOW_IMPORT, used by the doc build) `pipelines` -> `pipeline_wan_animate_2` -> `modular_pipelines.wan_animate_2` -> `modular_pipeline` -> `pipelines` is a cycle; importing the processor inside `__init__` breaks it. Goes away entirely with the standard pipeline's pre-merge removal. Co-Authored-By: Claude Fable 5 --- src/diffusers/pipelines/wan/pipeline_wan_animate_2.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/src/diffusers/pipelines/wan/pipeline_wan_animate_2.py b/src/diffusers/pipelines/wan/pipeline_wan_animate_2.py index f9b01c3dc63f..b4bc3ceb6c58 100644 --- a/src/diffusers/pipelines/wan/pipeline_wan_animate_2.py +++ b/src/diffusers/pipelines/wan/pipeline_wan_animate_2.py @@ -24,7 +24,6 @@ from ...loaders import WanLoraLoaderMixin from ...models import AutoencoderKLWan, WanAnimate2Transformer3DModel from ...models.transformers.transformer_wan_animate_2 import WanAnimate2KVCache -from ...modular_pipelines.wan_animate_2.video_processor import WanAnimate2VideoProcessor from ...schedulers import SchedulerMixin from ...utils import logging from ...utils.torch_utils import randn_tensor @@ -113,6 +112,11 @@ def __init__( ): super().__init__() + # Imported here rather than at module level: the modular package imports pipeline loading + # utilities from `pipelines`, so a module-level import back into `modular_pipelines` is a + # circular import under eager (DIFFUSERS_SLOW_IMPORT) initialization. + from ...modular_pipelines.wan_animate_2.video_processor import WanAnimate2VideoProcessor + self.register_modules( vae=vae, text_encoder=text_encoder, From 80938b276dbdb9be73c8b37df820a2baade361bf Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Wed, 12 Aug 2026 07:39:29 +0000 Subject: [PATCH 15/19] Default distilled sampling to 10 steps; fix docstring/signature mismatches - WanAnimate2DistilledCoreDenoiseStep overrides `inputs` to default num_inference_steps to 10 (the step count the distilled checkpoint is trained for), so the doc example needs no argument; base stays at 40 - transformer forward: give origin_area its own docstring entry and add the Returns: section; pipeline __call__: document callback_on_step_end, callback_on_step_end_tensor_inputs, max_sequence_length (fixes utils/check_forward_call_docstrings.py in CI) Co-Authored-By: Claude Fable 5 --- .../transformers/transformer_wan_animate_2.py | 13 ++++++++++--- .../modular_blocks_wan_animate_2_distilled.py | 14 +++++++++++--- .../pipelines/wan/pipeline_wan_animate_2.py | 7 +++++++ 3 files changed, 28 insertions(+), 6 deletions(-) diff --git a/src/diffusers/models/transformers/transformer_wan_animate_2.py b/src/diffusers/models/transformers/transformer_wan_animate_2.py index eef5cd168351..9259eae38342 100644 --- a/src/diffusers/models/transformers/transformer_wan_animate_2.py +++ b/src/diffusers/models/transformers/transformer_wan_animate_2.py @@ -801,13 +801,20 @@ def forward( 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*), origin_area (`list[int]`, *optional*): - Frame count and spatial size of the full video, which the in-context block mask is built over. 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}.") 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 index 660037342417..e01df0dd10c5 100644 --- 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 @@ -16,7 +16,7 @@ from ...utils import logging from ..modular_pipeline import SequentialPipelineBlocks -from ..modular_pipeline_utils import InsertableDict, OutputParam +from ..modular_pipeline_utils import InputParam, InsertableDict, OutputParam from .before_denoise import WanAnimate2PrepareSegmentsStep from .decoders import WanAnimate2DecodeStep from .denoise import WanAnimate2DistilledDenoiseStep @@ -184,7 +184,7 @@ class WanAnimate2DistilledCoreDenoiseStep(SequentialPipelineBlocks): 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): + 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 @@ -218,6 +218,14 @@ def description(self): "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 [ @@ -282,7 +290,7 @@ class WanAnimate2DistilledBlocks(SequentialPipelineBlocks): 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): + 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. diff --git a/src/diffusers/pipelines/wan/pipeline_wan_animate_2.py b/src/diffusers/pipelines/wan/pipeline_wan_animate_2.py index b4bc3ceb6c58..5b9c8f87beb1 100644 --- a/src/diffusers/pipelines/wan/pipeline_wan_animate_2.py +++ b/src/diffusers/pipelines/wan/pipeline_wan_animate_2.py @@ -300,6 +300,13 @@ def __call__( The output format. return_dict (`bool`, defaults to `True`): Whether to return a `WanPipelineOutput`. + callback_on_step_end (`Callable`, *optional*): + A function called at the end of each denoising step with `(self, step, timestep, callback_kwargs)`. + callback_on_step_end_tensor_inputs (`list[str]`, defaults to `["latents"]`): + The tensor inputs passed to `callback_on_step_end` through `callback_kwargs`; only tensors listed in + `self._callback_tensor_inputs` are allowed. + max_sequence_length (`int`, defaults to `512`): + Maximum sequence length for prompt encoding. """ # 1. Check inputs self.check_inputs(image, driving_video, prompt, height, width) From 3d1046ac7b60fb36233dd25e39555b037ac82b6f Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Wed, 12 Aug 2026 07:59:05 +0000 Subject: [PATCH 16/19] Pin Wan-Animate-2 block defaults; extend the defaults test for it MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit test_workflow_defaults now accepts the `None` key for pipelines without workflows (the full blockset as the single unnamed workflow) and an optional `component_configs` section pinning config values of `from_config` components against their creating spec — e.g. the guider scale that tells the two Wan-Animate-2 presets apart (3.0 vs 1.0). The Wan-Animate-2 testers pin components, the full input surface with defaults (40 vs 10 steps), and the guider scales. Co-Authored-By: Claude Fable 5 --- .../test_modular_pipelines_common.py | 14 +++++- .../test_modular_pipeline_wan_animate_2.py | 48 +++++++++++++++++++ 2 files changed, 61 insertions(+), 1 deletion(-) 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/test_modular_pipeline_wan_animate_2.py b/tests/modular_pipelines/wan_animate_2/test_modular_pipeline_wan_animate_2.py index 4b6048fdcb46..44e636de0e94 100644 --- 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 @@ -27,6 +27,52 @@ 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 @@ -36,6 +82,7 @@ class TestWanAnimate2ModularPipelineFast(ModularPipelineTesterMixin, ModularGuid 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) @@ -79,6 +126,7 @@ class TestWanAnimate2DistilledModularPipelineFast(TestWanAnimate2ModularPipeline 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): From 0eca946df0f042826d96a5966e11396362b2fe6a Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Wed, 12 Aug 2026 08:02:01 +0000 Subject: [PATCH 17/19] Document per-preset sampling defaults; testing guide for the defaults pin MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The Wan-Animate-2 API doc example needs no num_inference_steps anymore — each preset carries its own default (40 base, 10 distilled) and its own guidance. The testing guide documents expected_workflow_defaults: the None key for workflow-less blocksets and the optional component_configs pin for from_config components. Co-Authored-By: Claude Fable 5 --- .ai/testing.md | 2 +- docs/source/en/api/pipelines/wan_animate_2.md | 3 +-- 2 files changed, 2 insertions(+), 3 deletions(-) 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/api/pipelines/wan_animate_2.md b/docs/source/en/api/pipelines/wan_animate_2.md index e3b2ca3fcff3..7e23c8179ec3 100644 --- a/docs/source/en/api/pipelines/wan_animate_2.md +++ b/docs/source/en/api/pipelines/wan_animate_2.md @@ -50,13 +50,12 @@ videos = pipe( driving_video=driving_video, driving_video_fps=driving_video_fps, prompt="A cat in a blue uniform, white background", - num_inference_steps=40, 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 and pass `num_inference_steps=10`; no `guidance_scale` argument exists anywhere — guidance is owned by the pipeline's guider component. +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. From 612648ca7b21c73d6f6281f7d55539a9bc3028c6 Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Wed, 12 Aug 2026 08:17:46 +0000 Subject: [PATCH 18/19] Remove the standard WanAnimate2Pipeline The modular pipeline is the only Wan-Animate-2 entry point, per the plan (kept during review only) and with the author's ok. Every refactor of the modular decomposition was verified bit-identical against this pipeline while it existed. The clip_visual_encode / get_i2v_mask / get_frame_indices helpers in the modular folder become canonical (their "Copied from" sources are gone), and the circular import the pipeline's processor import created disappears with it. Co-Authored-By: Claude Fable 5 --- src/diffusers/__init__.py | 2 - .../wan_animate_2/encoders.py | 3 - src/diffusers/pipelines/__init__.py | 2 - src/diffusers/pipelines/wan/__init__.py | 2 - .../pipelines/wan/pipeline_wan_animate_2.py | 642 ------------------ .../dummy_torch_and_transformers_objects.py | 15 - 6 files changed, 666 deletions(-) delete mode 100644 src/diffusers/pipelines/wan/pipeline_wan_animate_2.py diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py index f0d608a27e02..66cfa442908b 100644 --- a/src/diffusers/__init__.py +++ b/src/diffusers/__init__.py @@ -854,7 +854,6 @@ "VisualClozeGenerationPipeline", "VisualClozePipeline", "VQDiffusionPipeline", - "WanAnimate2Pipeline", "WanAnimatePipeline", "WanImageToVideoPipeline", "WanPipeline", @@ -1686,7 +1685,6 @@ VisualClozeGenerationPipeline, VisualClozePipeline, VQDiffusionPipeline, - WanAnimate2Pipeline, WanAnimatePipeline, WanImageToVideoPipeline, WanPipeline, diff --git a/src/diffusers/modular_pipelines/wan_animate_2/encoders.py b/src/diffusers/modular_pipelines/wan_animate_2/encoders.py index 4100c9aa0181..21b70f636f7d 100644 --- a/src/diffusers/modular_pipelines/wan_animate_2/encoders.py +++ b/src/diffusers/modular_pipelines/wan_animate_2/encoders.py @@ -68,7 +68,6 @@ def get_t5_prompt_embeds( return prompt_embeds -# Copied from diffusers.pipelines.wan.pipeline_wan_animate_2.clip_visual_encode def clip_visual_encode(image_encoder, tensor, device, dtype): """Encode tensor to CLIP features (bicubic to 224×224, matching original).""" if tensor.ndim == 3: @@ -82,7 +81,6 @@ def clip_visual_encode(image_encoder, tensor, device, dtype): return out.hidden_states[-2] -# Copied from diffusers.pipelines.wan.pipeline_wan_animate_2.get_i2v_mask def get_i2v_mask(lat_t, lat_h, lat_w, mask_len=1, device="cuda"): """Create an i2v mask in latent space. @@ -96,7 +94,6 @@ def get_i2v_mask(lat_t, lat_h, lat_w, mask_len=1, device="cuda"): return msk -# Copied from diffusers.pipelines.wan.pipeline_wan_animate_2.get_frame_indices 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) diff --git a/src/diffusers/pipelines/__init__.py b/src/diffusers/pipelines/__init__.py index 02f7a478dd8c..50052b0ca887 100644 --- a/src/diffusers/pipelines/__init__.py +++ b/src/diffusers/pipelines/__init__.py @@ -435,7 +435,6 @@ "WanVideoToVideoPipeline", "WanVACEPipeline", "WanAnimatePipeline", - "WanAnimate2Pipeline", ] _import_structure["kandinsky5"] = [ "Kandinsky5T2VPipeline", @@ -898,7 +897,6 @@ ) from .visualcloze import VisualClozeGenerationPipeline, VisualClozePipeline from .wan import ( - WanAnimate2Pipeline, WanAnimatePipeline, WanImageToVideoPipeline, WanPipeline, diff --git a/src/diffusers/pipelines/wan/__init__.py b/src/diffusers/pipelines/wan/__init__.py index 3eac9b5a666b..ad51a52f9242 100644 --- a/src/diffusers/pipelines/wan/__init__.py +++ b/src/diffusers/pipelines/wan/__init__.py @@ -24,7 +24,6 @@ else: _import_structure["pipeline_wan"] = ["WanPipeline"] _import_structure["pipeline_wan_animate"] = ["WanAnimatePipeline"] - _import_structure["pipeline_wan_animate_2"] = ["WanAnimate2Pipeline"] _import_structure["pipeline_wan_i2v"] = ["WanImageToVideoPipeline"] _import_structure["pipeline_wan_vace"] = ["WanVACEPipeline"] _import_structure["pipeline_wan_video2video"] = ["WanVideoToVideoPipeline"] @@ -38,7 +37,6 @@ else: from .pipeline_wan import WanPipeline from .pipeline_wan_animate import WanAnimatePipeline - from .pipeline_wan_animate_2 import WanAnimate2Pipeline from .pipeline_wan_i2v import WanImageToVideoPipeline from .pipeline_wan_vace import WanVACEPipeline from .pipeline_wan_video2video import WanVideoToVideoPipeline diff --git a/src/diffusers/pipelines/wan/pipeline_wan_animate_2.py b/src/diffusers/pipelines/wan/pipeline_wan_animate_2.py deleted file mode 100644 index 5b9c8f87beb1..000000000000 --- a/src/diffusers/pipelines/wan/pipeline_wan_animate_2.py +++ /dev/null @@ -1,642 +0,0 @@ -# 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 -from typing import Any, Callable - -import numpy as np -import torch -import torch.nn.functional as F -from tqdm import tqdm - -from ...image_processor import PipelineImageInput -from ...loaders import WanLoraLoaderMixin -from ...models import AutoencoderKLWan, WanAnimate2Transformer3DModel -from ...models.transformers.transformer_wan_animate_2 import WanAnimate2KVCache -from ...schedulers import SchedulerMixin -from ...utils import logging -from ...utils.torch_utils import randn_tensor -from ..pipeline_utils import DiffusionPipeline -from .pipeline_output import WanPipelineOutput - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - - -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 - - -CLIP_MEAN = [0.48145466, 0.4578275, 0.40821073] -CLIP_STD = [0.26862954, 0.26130258, 0.27577711] - - -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 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] - - -class WanAnimate2Pipeline(DiffusionPipeline, WanLoraLoaderMixin): - r""" - Pipeline for character animation using Wan-Animate-2. - - This pipeline takes a reference character image and a driving video, and generates a video where the character is - animated following the motion in the driving video. The model uses an in-context attention mechanism with KV cache: - a reference video is first encoded to cache K/V tensors, then the generation forward uses the cached K/V with a - block mask for frame-level sparse in-context attention. - - Args: - tokenizer ([`AutoTokenizer`]): - Tokenizer for the umT5 text encoder. - text_encoder ([`UMT5EncoderModel`]): - The umT5 text encoder. - image_encoder ([`CLIPVisionModel`]): - CLIP vision model for encoding the reference image. - transformer ([`WanAnimate2Transformer3DModel`]): - The Wan-Animate-2 transformer model. - scheduler ([`SchedulerMixin`]): - A flow-matching scheduler to be used in combination with `transformer` to denoise the encoded latents. The - reference implementation samples with `DPMSolverMultistepScheduler` (`flow_shift=5.0`) for the base model - and `FlowMatchEulerDiscreteScheduler` (`shift=5.0`) for the distilled one. - vae ([`AutoencoderKLWan`]): - The Wan VAE model. - """ - - model_cpu_offload_seq = "text_encoder->image_encoder->transformer->vae" - _callback_tensor_inputs = ["latents"] - - def __init__( - self, - tokenizer, - text_encoder, - vae: AutoencoderKLWan, - scheduler: SchedulerMixin, - image_encoder, - transformer: WanAnimate2Transformer3DModel, - ): - super().__init__() - - # Imported here rather than at module level: the modular package imports pipeline loading - # utilities from `pipelines`, so a module-level import back into `modular_pipelines` is a - # circular import under eager (DIFFUSERS_SLOW_IMPORT) initialization. - from ...modular_pipelines.wan_animate_2.video_processor import WanAnimate2VideoProcessor - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - tokenizer=tokenizer, - image_encoder=image_encoder, - transformer=transformer, - scheduler=scheduler, - ) - - self.vae_scale_factor_temporal = self.vae.config.scale_factor_temporal if getattr(self, "vae", None) else 4 - self.vae_scale_factor_spatial = self.vae.config.scale_factor_spatial if getattr(self, "vae", None) else 8 - # Wan-Animate-2 letterboxes the reference image and the driving video into the same frame: aspect - # ratio preserved, the remainder filled with black (`resize_mode="fill"` with `fill_color=0`). - # The reference implementation resizes with cv2, which is not a diffusers dependency, so these - # processors use the closest PIL kernels: bilinear for the driving frames (the same filter as - # `INTER_LINEAR`, but PIL quantizes interpolation weights to 22 bits where cv2 uses 11, so - # exact-half values round in opposite directions -- at most one 8-bit level per pixel) and - # bicubic for the reference image's downscale (`INTER_AREA` has no PIL equivalent; bicubic - # measures closest). Outputs therefore differ very slightly, numerically and visually, from - # the original repository. - self.image_processor_for_reference = WanAnimate2VideoProcessor( - vae_scale_factor=self.vae_scale_factor_spatial, spatial_patch_size=(2, 2), resample="bicubic" - ) - self.video_processor = WanAnimate2VideoProcessor( - vae_scale_factor=self.vae_scale_factor_spatial, spatial_patch_size=(2, 2), resample="bilinear" - ) - - def _get_t5_prompt_embeds(self, prompt, device=None, dtype=None, max_sequence_length=512): - device = device or self._execution_device - dtype = dtype or self.text_encoder.dtype - - prompt = [prompt] if isinstance(prompt, str) else prompt - - text_inputs = self.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 = self.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 encode_image(self, image, device=None): - device = device or self._execution_device - from transformers import CLIPImageProcessor - - image_processor = CLIPImageProcessor() - processed = image_processor(images=image, return_tensors="pt").to(device) - image_embeds = self.image_encoder(**processed, output_hidden_states=True) - return image_embeds.hidden_states[-2] - - def _encode_vae(self, video): - """VAE-encode a `[B, C, T, H, W]` clip and standardize the latents.""" - latents = self.vae.encode(video.to(self.vae.dtype)) - if hasattr(latents, "latent_dist"): - latents = latents.latent_dist.mode() - elif hasattr(latents, "latents"): - latents = latents.latents - elif isinstance(latents, (list, tuple)): - latents = latents[0] if isinstance(latents[0], torch.Tensor) else torch.stack(latents) - latents_mean = ( - torch.tensor(self.vae.config.latents_mean) - .view(1, self.vae.config.z_dim, 1, 1, 1) - .to(latents.device, latents.dtype) - ) - latents_recip_std = 1.0 / torch.tensor(self.vae.config.latents_std).view(1, self.vae.config.z_dim, 1, 1, 1).to( - latents.device, latents.dtype - ) - latents = (latents - latents_mean) * latents_recip_std - return latents - - def _decode_vae(self, latents, device): - """Decode latents to video using VAE, with destandardization.""" - latents = latents.to(self.vae.dtype) - latents_mean = ( - torch.tensor(self.vae.config.latents_mean) - .view(1, self.vae.config.z_dim, 1, 1, 1) - .to(latents.device, latents.dtype) - ) - latents_recip_std = 1.0 / torch.tensor(self.vae.config.latents_std).view(1, self.vae.config.z_dim, 1, 1, 1).to( - latents.device, latents.dtype - ) - latents = latents / latents_recip_std + latents_mean - out_frames = self.vae.decode(latents, return_dict=False)[0] - return out_frames - - def check_inputs(self, image, driving_video, prompt, height, width): - if image is None: - raise ValueError("Provide `image`. Cannot leave `image` undefined.") - if driving_video is None: - raise ValueError("Provide `driving_video`. Cannot leave `driving_video` undefined.") - if height % 16 != 0 or width % 16 != 0: - raise ValueError(f"`height` and `width` have to be divisible by 16 but are {height} and {width}.") - - @property - def guidance_scale(self): - return self._guidance_scale - - @property - def do_classifier_free_guidance(self): - return self._guidance_scale > 1 - - @property - def num_timesteps(self): - return self._num_timesteps - - @torch.no_grad() - def __call__( - self, - image: PipelineImageInput, - driving_video: list[Any], - prompt: str | list[str] = None, - negative_prompt: str | list[str] = None, - prompt_ref: str = "人物动作的参考视频", - height: int = 800, - width: int = 640, - segment_frame_length: int = 81, - prev_segment_conditioning_frames: int = 1, - fps: int = 24, - driving_video_fps: float | None = None, - num_inference_steps: int = 40, - guidance_scale: float = 3.0, - generator: torch.Generator | list[torch.Generator] | None = None, - output_type: str | None = "np", - return_dict: bool = True, - callback_on_step_end: Callable | None = None, - callback_on_step_end_tensor_inputs: list[str] = ["latents"], - max_sequence_length: int = 512, - ): - r""" - The call function for character animation generation. - - Args: - image (`PipelineImageInput`): - The reference character image. - driving_video (`list[PIL.Image.Image]`, `np.ndarray` or `torch.Tensor`): - The driving video that provides the motion, in any format accepted by - [`~video_processor.VideoProcessor.preprocess_video`]. Load one from disk with [`~utils.load_video`]. - prompt (`str` or `list[str]`): - The text prompt describing the character appearance and background. - negative_prompt (`str` or `list[str]`, *optional*): - The negative prompt for classifier-free guidance. - prompt_ref (`str`, defaults to `"人物动作的参考视频"`): - The reference prompt for the driving video context. - height (`int`, defaults to `800`): - Together with `width`, the target *area* (`height * width`) of the generated video. The aspect ratio is - taken from `image`, so the video is rarely exactly `height` x `width` — both dimensions are rescaled to - hit that area and then floored to a multiple of 16. - width (`int`, defaults to `640`): - See `height`. - segment_frame_length (`int`, defaults to `81`): - The number of frames in each inference segment. - prev_segment_conditioning_frames (`int`, defaults to `1`): - The number of conditioning frames from the previous segment. - fps (`int`, defaults to `24`): - The frame rate the model generates at. `driving_video` is resampled to it when `driving_video_fps` is - given. - driving_video_fps (`float`, *optional*): - The frame rate `driving_video` was captured at — a list of frames does not carry it, so - [`~utils.load_video`] will report it with `return_fps=True`. When set, the driving frames are - nearest-neighbour resampled from it to `fps`; when `None` they are used as-is. - num_inference_steps (`int`, defaults to `40`): - The number of denoising steps. - guidance_scale (`float`, defaults to `3.0`): - Guidance scale for classifier-free guidance. - generator (`torch.Generator`, *optional*): - A generator to make generation deterministic. - output_type (`str`, defaults to `"np"`): - The output format. - return_dict (`bool`, defaults to `True`): - Whether to return a `WanPipelineOutput`. - callback_on_step_end (`Callable`, *optional*): - A function called at the end of each denoising step with `(self, step, timestep, callback_kwargs)`. - callback_on_step_end_tensor_inputs (`list[str]`, defaults to `["latents"]`): - The tensor inputs passed to `callback_on_step_end` through `callback_kwargs`; only tensors listed in - `self._callback_tensor_inputs` are allowed. - max_sequence_length (`int`, defaults to `512`): - Maximum sequence length for prompt encoding. - """ - # 1. Check inputs - self.check_inputs(image, driving_video, prompt, height, width) - - self._guidance_scale = guidance_scale - device = self._execution_device - - # 2. Resolve the output frame. `height * width` is a target *area*; the aspect ratio comes from the - # reference image, and both sides are floored to a multiple of `vae_scale_factor_spatial * patch_size` - # so the latent grid divides evenly. - image_height, image_width = self.video_processor.get_default_height_width(image) - mod_value = self.vae_scale_factor_spatial * 2 - aspect_ratio = image_height / image_width - actual_h = int(math.sqrt(height * width * aspect_ratio)) // mod_value * mod_value - actual_w = int(math.sqrt(height * width / aspect_ratio)) // mod_value * mod_value - - # The reference image is letterboxed into that frame. `resize_mode="fill"` keeps the aspect ratio and - # pads the remainder with black; record the pasted box so the bars can be cropped back off the output. - src_w = ( - actual_w if actual_w / actual_h < image_width / image_height else image_width * actual_h // image_height - ) - src_h = ( - actual_h if actual_w / actual_h >= image_width / image_height else image_height * actual_w // image_width - ) - crop_top, crop_left = (actual_h - src_h) // 2, (actual_w - src_w) // 2 - - image_pixels = self.image_processor_for_reference.preprocess( - image, height=actual_h, width=actual_w, resize_mode="fill" - ).to(device, dtype=torch.float32) - - # 3. Preprocess the driving video into the same frame, resampling to `fps` first if asked to. - if driving_video_fps is not None: - frame_indices = get_frame_indices(len(driving_video), driving_video_fps, fps) - driving_video = [driving_video[i] for i in frame_indices] - - driving_video = self.video_processor.preprocess_video( - driving_video, height=actual_h, width=actual_w, resize_mode="fill" - ).to(device, dtype=torch.float32) - - # Pad driving video to be a multiple of (segment_frame_length - prev_segment_conditioning_frames) - real_frame_len = driving_video.shape[2] - effective_segment = segment_frame_length - prev_segment_conditioning_frames - last_segment_frames = ( - (real_frame_len - prev_segment_conditioning_frames) % effective_segment - if real_frame_len > prev_segment_conditioning_frames - else 0 - ) - if last_segment_frames > 0: - num_padding = effective_segment - last_segment_frames - else: - num_padding = 0 - target_num_frames = real_frame_len + num_padding - - # Pad driving video using zigzag (reflect) strategy - if num_padding > 0: - # Mirrored real frames, not filler: the model attends to them like any other frame and needs no mask. - # The surplus generated frames are cropped off again with `[:, :, :real_frame_len]` at the end. - padding_frames = driving_video[:, :, real_frame_len - num_padding : real_frame_len].flip(2) - driving_video = torch.cat([driving_video, padding_frames], dim=2) - - # 4. Encode prompt - prompt_embeds = self._get_t5_prompt_embeds(prompt, device=device, max_sequence_length=max_sequence_length) - negative_prompt_embeds = None - if self.do_classifier_free_guidance: - negative_prompt = negative_prompt or "" - negative_prompt_embeds = self._get_t5_prompt_embeds( - negative_prompt, device=device, max_sequence_length=max_sequence_length - ) - - # Reference prompt - prompt_ref_embeds = self._get_t5_prompt_embeds( - prompt_ref, device=device, max_sequence_length=max_sequence_length - ) - - # 5. Encode reference image (VAE + CLIP) - # CLIP features from reference image (direct bicubic to 224×224 from tensor) - clip_fea = clip_visual_encode(self.image_encoder, image_pixels[0], device, self.transformer.dtype) - - ref_latents = self._encode_vae(image_pixels.unsqueeze(2)) # [B, C, H, W] -> [B, C, 1, H, W] - - latent_h = actual_h // self.vae_scale_factor_spatial - latent_w = actual_w // self.vae_scale_factor_spatial - - # Prepare reference i2v mask and y_ref - mask_ref = get_i2v_mask(1, latent_h, latent_w, 1, device=device).to(self.transformer.dtype) - ref_lat_0 = ref_latents[0] if ref_latents.ndim == 5 else ref_latents - y_ref = torch.cat([mask_ref, ref_lat_0], dim=0) - - # CLIP context for reference - clip_context = clip_fea - - # 6. Prepare timesteps - self.scheduler.set_timesteps(num_inference_steps, device=device) - timesteps = self.scheduler.timesteps - self._num_timesteps = len(timesteps) - - # 7. Segment-based generation loop - start = 0 - end = segment_frame_length - all_out_frames = [] - out_frames = None - - num_segments = ( - target_num_frames - prev_segment_conditioning_frames + effective_segment - 1 - ) // effective_segment - - for seg_idx in range(num_segments): - if start + prev_segment_conditioning_frames >= target_num_frames: - break - - mask_reft_len = prev_segment_conditioning_frames if start > 0 else 0 - - if target_num_frames - start < segment_frame_length: - segment_frame_length_actual = target_num_frames - start - else: - segment_frame_length_actual = segment_frame_length - - # VAE-encode this segment's slice of the driving video. The Wan VAE is causal in time, so - # encoding the whole video once up front and slicing the latents is not the same tensor — - # segments overlap by `prev_segment_conditioning_frames` frames and each slice restarts the temporal convolution. - # Encoding per segment is also what a streaming mode would have to do anyway. - condition_latents = self._encode_vae(driving_video[:, :, start : start + segment_frame_length_actual]) - - # CLIP features from driving video first frame (direct bicubic to 224×224 from tensor) - condition_img = driving_video[0, :, 0] # [C, H, W] in [-1, 1] - condition_clip_context = clip_visual_encode( - self.image_encoder, condition_img, device, self.transformer.dtype - ) - - # Prepare condition y (mask + latents) - T = segment_frame_length_actual + 1 - - # Encode condition y - if mask_reft_len > 0: - prev_frames = out_frames[0, :, -mask_reft_len:].clone().detach() - prev_frames_interp = F.interpolate( - prev_frames.permute(1, 0, 2, 3), size=(actual_h, actual_w), mode="bicubic" - ).permute(1, 0, 2, 3) - cond_y_input = torch.cat( - [prev_frames_interp, torch.zeros(3, T - mask_reft_len - 1, actual_h, actual_w, device=device)], - dim=1, - ) - else: - cond_y_input = torch.zeros(3, T - 1, actual_h, actual_w, device=device) - - y_reft = self._encode_vae(cond_y_input.unsqueeze(0)) - if y_reft.ndim == 5: - y_reft = y_reft.squeeze(0) # [1, 16, T, H, W] -> [16, T, H, W] - - # Derive lat_t from actual VAE output shape - lat_t_y = y_reft.shape[1] # temporal dimension of y_reft latents - lat_t_cond = condition_latents.shape[2] if condition_latents.ndim == 5 else condition_latents.shape[1] - - msk_reft = get_i2v_mask(lat_t_y, latent_h, latent_w, mask_reft_len, device=device).to( - self.transformer.dtype - ) - y_reft = torch.cat([msk_reft, y_reft], dim=0) - - # Condition mask and latents - condition_msk_y = get_i2v_mask( - lat_t_cond, latent_h, latent_w, segment_frame_length_actual, device=device - ).to(self.transformer.dtype) - cond_lat_0 = condition_latents[0] if condition_latents.ndim == 5 else condition_latents - condition_y = torch.cat([condition_msk_y, cond_lat_0], dim=0) - - y = torch.cat([y_ref, y_reft], dim=1) - - # Prepare grid sizes — use post-patch spatial dims (VAE 8x + patch 2x = 16x total) - if condition_latents.ndim == 5: - ref_shape = list(condition_latents.shape[2:]) # [T, H, W] pre-patch - else: - ref_shape = list(condition_latents.shape[1:]) - # After patch_embedding (1,2,2): spatial dims halved - ref_shape_post = [ref_shape[0], ref_shape[1] // 2, ref_shape[2] // 2] - grid_sizes_ref = torch.tensor([ref_shape_post], dtype=torch.long) - - # Noise latents temporal dim = y_ref(1) + y_reft/condition_y(T) = total y temporal dim - lat_t_noise = y.shape[1] if y.ndim == 4 else y.shape[2] - noise = randn_tensor( - (16, lat_t_noise, latent_h, latent_w), - generator=generator, - device=device, - dtype=torch.float32, - ) - - latents = [noise] - - # Prepare arguments for transformer - max_seq_len = int(math.ceil(np.prod([lat_t_noise, latent_h // 2, latent_w // 2]))) - max_seq_len_ref = int(math.ceil(np.prod(ref_shape) // 4)) if ref_shape else max_seq_len - - arg_c = { - "context": [prompt_embeds[0]], - "seq_len": max_seq_len, - "clip_fea": clip_context, - "y": [y], - "origin_len": segment_frame_length_actual, - "origin_area": [actual_h, actual_w], - } - - arg_ref_c = { - "context_ref": [prompt_ref_embeds[0]], - "seq_len_ref": max_seq_len_ref, - "clip_fea_ref": condition_clip_context, - "y_ref": [condition_y], - } - - arg_null = None - if self.do_classifier_free_guidance: - arg_null = { - "context": [negative_prompt_embeds[0]], - "seq_len": max_seq_len, - "clip_fea": clip_context, - "y": [y], - "origin_len": segment_frame_length_actual, - "origin_area": [actual_h, actual_w], - "is_uncondtion": True, - } - - kv_cache = WanAnimate2KVCache(self.transformer.config.num_layers) - - # Phase 1: encode reference — cast all inputs to transformer dtype - t_ref = torch.tensor([timesteps[0].item()], device=device, dtype=self.transformer.dtype) - self.transformer( - [condition_latents[0].to(self.transformer.dtype)] - if condition_latents.ndim == 5 - else [condition_latents.to(self.transformer.dtype)], - timestep=t_ref, - encoder_hidden_states=[c.to(self.transformer.dtype) for c in arg_ref_c["context_ref"]], - encoder_hidden_states_image=arg_ref_c["clip_fea_ref"].to(self.transformer.dtype), - condition_latents=[y.to(self.transformer.dtype) for y in arg_ref_c["y_ref"]], - kv_cache=kv_cache, - kv_cache_mode="extract", - seq_len=max_seq_len_ref, - offset_grid_sizes=grid_sizes_ref, - ) - - # Phase 2: denoising loop - for i, t in tqdm(enumerate(timesteps), total=len(timesteps), desc=f"Segment {seg_idx + 1}/{num_segments}"): - timestep = torch.stack([t]) - - # Conditional - noise_pred_cond = self.transformer( - [l.to(self.transformer.dtype) for l in latents], - timestep=timestep, - encoder_hidden_states=arg_c["context"], - encoder_hidden_states_image=arg_c["clip_fea"], - condition_latents=arg_c["y"], - kv_cache=kv_cache, - kv_cache_mode="cached", - seq_len=max_seq_len, - reference_grid_sizes=grid_sizes_ref, - origin_len=arg_c["origin_len"], - origin_area=arg_c["origin_area"], - ).sample[0] - - if self.do_classifier_free_guidance: - noise_pred_uncond = self.transformer( - [l.to(self.transformer.dtype) for l in latents], - timestep=timestep, - encoder_hidden_states=arg_null["context"], - encoder_hidden_states_image=arg_null["clip_fea"], - condition_latents=arg_null["y"], - kv_cache=kv_cache, - kv_cache_mode="cached", - seq_len=max_seq_len, - reference_grid_sizes=grid_sizes_ref, - origin_len=arg_null["origin_len"], - origin_area=arg_null["origin_area"], - is_uncondtion=True, - ).sample[0] - - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_cond - noise_pred_uncond) - else: - noise_pred = noise_pred_cond - - # Scheduler step - temp_x0 = self.scheduler.step( - noise_pred.unsqueeze(0), - t, - latents[0].unsqueeze(0), - return_dict=False, - generator=generator, - )[0] - latents[0] = temp_x0.squeeze(0) - - if callback_on_step_end is not None: - callback_kwargs = {} - for k in callback_on_step_end_tensor_inputs: - callback_kwargs[k] = locals()[k] - callback_outputs = callback_on_step_end(self, i, t, callback_kwargs) - latents[0] = ( - callback_outputs.pop("latents", latents)[0] - if isinstance(callback_outputs.get("latents"), list) - else latents[0] - ) - - # Decode - x0 = [latents[0].to(dtype=torch.float32)] - out_frames = self._decode_vae(x0[0][:, 1:], device) - - if start > 0: - out_frames = out_frames[:, :, mask_reft_len:] - - all_out_frames.append(out_frames.cpu()) - start += effective_segment - end += effective_segment - - # Each segment allocates a fresh KV cache — at 720p that is tens of GB, and holding - # the previous one while the next is built fragments the allocator enough to OOM. - kv_cache.clear() - # `out_frames` is deliberately kept: the next segment conditions on its tail. - del kv_cache, latents, x0 - torch.cuda.empty_cache() - - # Each segment is an independent trajectory, so the solver state has to be reset. - self.scheduler.set_timesteps(num_inference_steps, device=device) - timesteps = self.scheduler.timesteps - - # Concatenate all segments - video = torch.cat(all_out_frames, dim=2)[:, :, :real_frame_len].to(device) - - # Crop the reference image's letterbox bars back off - video = video[:, :, :, crop_top : crop_top + src_h, crop_left : crop_left + src_w] - - video = self.video_processor.postprocess_video(video, output_type=output_type) - - self.maybe_free_model_hooks() - - if not return_dict: - return (video,) - - return WanPipelineOutput(frames=video) diff --git a/src/diffusers/utils/dummy_torch_and_transformers_objects.py b/src/diffusers/utils/dummy_torch_and_transformers_objects.py index 3301b211148a..11e69fcdccf1 100644 --- a/src/diffusers/utils/dummy_torch_and_transformers_objects.py +++ b/src/diffusers/utils/dummy_torch_and_transformers_objects.py @@ -5147,21 +5147,6 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch", "transformers"]) -class WanAnimate2Pipeline(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 WanAnimatePipeline(metaclass=DummyObject): _backends = ["torch", "transformers"] From e38e13d5fdbb529b52346801acf8a3c2f6f25ea3 Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Wed, 12 Aug 2026 08:28:09 +0000 Subject: [PATCH 19/19] Give the distilled Wan-Animate-2 preset its own model name MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit wan-animate-2-distilled maps to WanAnimate2DistilledModularPipeline, and every distilled block class carries it, so distilled blocks init (and save_pretrained round-trip as) the distilled pipeline class instead of the base one. Shared leaves keep wan-animate-2 — standalone they route to the base class, which only differs in its default blocks. Co-Authored-By: Claude Fable 5 --- src/diffusers/modular_pipelines/modular_pipeline.py | 1 + src/diffusers/modular_pipelines/wan_animate_2/denoise.py | 4 ++++ .../modular_blocks_wan_animate_2_distilled.py | 8 ++++---- 3 files changed, 9 insertions(+), 4 deletions(-) diff --git a/src/diffusers/modular_pipelines/modular_pipeline.py b/src/diffusers/modular_pipelines/modular_pipeline.py index d865beb066d9..88cea1b78b1d 100644 --- a/src/diffusers/modular_pipelines/modular_pipeline.py +++ b/src/diffusers/modular_pipelines/modular_pipeline.py @@ -133,6 +133,7 @@ def _helios_pyramid_map_fn(config_dict=None): ("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/denoise.py b/src/diffusers/modular_pipelines/wan_animate_2/denoise.py index 7a97c1c13db5..d96b8f814239 100644 --- a/src/diffusers/modular_pipelines/wan_animate_2/denoise.py +++ b/src/diffusers/modular_pipelines/wan_animate_2/denoise.py @@ -590,6 +590,8 @@ def __call__(self, components, block_state: BlockState, k: int): class WanAnimate2DistilledSegmentDenoiseInner(WanAnimate2SegmentDenoiseInner): + model_name = "wan-animate-2-distilled" + @property def description(self) -> str: return ( @@ -779,6 +781,8 @@ def description(self) -> str: class WanAnimate2DistilledDenoiseStep(WanAnimate2SegmentLoopWrapper): + model_name = "wan-animate-2-distilled" + block_classes = [ WanAnimate2SegmentVaeEncoderStep, WanAnimate2SegmentPrevFramesStep, 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 index e01df0dd10c5..8eab815897da 100644 --- 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 @@ -73,7 +73,7 @@ class WanAnimate2DistilledImageEncodeStep(SequentialPipelineBlocks): CLIP vision features of the reference image, conditioning every denoising forward """ - model_name = "wan-animate-2" + model_name = "wan-animate-2-distilled" block_classes = WanAnimate2DistilledImageEncoderBlocks.values() block_names = WanAnimate2DistilledImageEncoderBlocks.keys() @@ -133,7 +133,7 @@ class WanAnimate2DistilledVideoEncodeStep(SequentialPipelineBlocks): CLIP vision features of the driving video's first frame """ - model_name = "wan-animate-2" + model_name = "wan-animate-2-distilled" block_classes = WanAnimate2DistilledVideoEncoderBlocks.values() block_names = WanAnimate2DistilledVideoEncoderBlocks.keys() @@ -207,7 +207,7 @@ class WanAnimate2DistilledCoreDenoiseStep(SequentialPipelineBlocks): them into the final video """ - model_name = "wan-animate-2" + model_name = "wan-animate-2-distilled" block_classes = WanAnimate2DistilledCoreDenoiseBlocks.values() block_names = WanAnimate2DistilledCoreDenoiseBlocks.keys() @@ -302,7 +302,7 @@ class WanAnimate2DistilledBlocks(SequentialPipelineBlocks): The generated videos. """ - model_name = "wan-animate-2" + model_name = "wan-animate-2-distilled" block_classes = DISTILLED_BLOCKS.values() block_names = DISTILLED_BLOCKS.keys()