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