Diffusers 中 WanAnimateTransformer3DModel 深度解析:Wan2.2 面部动画驱动 Diffusion Transformer 的架构与使用
【免费下载链接】diffusers🤗 Diffusers: State-of-the-art diffusion models for image, video, and audio generation in PyTorch.项目地址: https://gitcode.com/GitHub_Trending/di/diffusers
导读
WanAnimateTransformer3DModel是 🤗 Diffusers 仓库中为 Wan Animate(阿里 Wan 团队 Wan2.2 系列)实现的面部动画驱动扩散 Transformer 模型,它以 3D 视频潜变量为处理对象,将参考人物图像、姿态视频、面部视频与文本提示统一融合,实现高质量的角色动画生成。本文以 WanAnimateTransformer3DModel API 文档 为骨架,结合 transformer_wan_animate.py 源码与 pipeline_wan_animate.py 调用链,带你掌握模型的加载方式、完整架构、全部配置参数、前向输入输出约定,以及它在 WanAnimatePipeline 中的实际工作方式。
模型定位:为"面部驱动动画"而生的扩散 Transformer
WanAnimateTransformer3DModel 在 Wan Animate API 文档 中定位为"用于类视频 3D 数据的 Diffusion Transformer 模型",由阿里 Wan 团队在 Wan2.2 系列中提出。与普通文生视频 DiT 不同,它的核心任务是以人物面部与姿态的运动信息作为条件信号驱动生成,因此模型内部集成了两套为动画任务专门设计的编码器:
- 运动编码器(Motion Encoder):从面部视频的像素序列中提取紧凑的运动向量;
- 面部编码器(Face Encoder):把逐帧运动向量转化为可与 Transformer 主分支对齐的时间对齐特征,并通过"面部适配器"注入到主干网络。
从源码类定义(transformer_wan_animate.py#L986-L988)可以看到,该类同时继承了ModelMixin、ConfigMixin、PeftAdapterMixin、FromOriginalModelMixin、CacheMixin与AttentionMixin,这意味着它天然支持from_pretrained加载、LoRA/Peft 适配、原始权重转换、条件缓存与自定义注意力后端等 Diffusers 标准能力。
快速加载:一行代码加载 14B 动画 Transformer
官方文档给出的加载方式如下:
from diffusers import WanAnimateTransformer3DModel transformer = WanAnimateTransformer3DModel.from_pretrained( "Wan-AI/Wan2.2-Animate-14B-Diffusers", subfolder="transformer", dtype=torch.bfloat16 )WanAnimateTransformer3DModel已通过 models/init.py#L153 注册进diffusers顶层命名空间,因此可以直接from diffusers import。为了让上面的代码真正可运行,需要补全torch导入并指定设备:
import torch from diffusers import WanAnimateTransformer3DModel transformer = WanAnimateTransformer3DModel.from_pretrained( "Wan-AI/Wan2.2-Animate-14B-Diffusers", subfolder="transformer", torch_dtype=torch.bfloat16, ) transformer.to("cuda")几个值得注意的加载细节:
subfolder="transformer":仓库目录内模型权重存放于transformer子目录,与 VAE、文本编码器(umT5、CLIP)等组件分开管理;dtype=torch.bfloat16:14B 级别参数在 FP32 下显存压力极大,官方推荐以 bf16 加载推理;- 从原始权重加载:由于类实现了
FromOriginalModelMixin,也支持通过 single-file 方式从.gguf等单文件权重恢复,详见下文"量化与单文件加载"一节。
架构全景:六段式前向流水线
forward的实现(transformer_wan_animate.py#L1155-L1306)清晰地划分为六个阶段,注释与官方实现一一对应:
- 旋转位置编码(RoPE):
WanRotaryPosEmbed按 T/H/W 三个维度分别生成 1D 旋转位置频率并拼接,输出形状为(1, ppf * pph * ppw, 1, head_dim); - Patch 嵌入(Patch Embedding):
patch_embedding与pose_patch_embedding是两个Conv3d,前者把含参考帧的视频潜变量(B, 2C+4, T+1, H, W)打成 patch,后者单独嵌入姿态潜变量,随后姿态嵌入被加到主分支的后 T 帧上(hidden_states[:, :, 1:] += pose_hidden_states,见 L1223),再flatten(2).transpose(1, 2)变成序列形式; - 条件嵌入(时间/文本/图像):
WanTimeTextImageEmbedding基于 Wan2.1 的时间步逻辑,输出temb、timestep_proj以及处理后的文本/图像嵌入,图像嵌入(CLIP 视觉特征)会被concat到文本嵌入前方形成联合条件; - 运动特征提取:面部视频像素先经
WanAnimateMotionEncoder得到逐帧运动向量,再由WanAnimateFaceEncoder转成带时间因果性的面部特征; - Transformer 主干 + 面部适配器:40 层
WanTransformerBlock顺序执行,每 5 层(inject_face_latents_blocks)插入一次WanAnimateFaceBlockCrossAttention的交叉注意力结果(残差相加); - 输出归一化与反 patchify:
norm_out(FP32 LayerNorm)加 scale-shift 调制后经proj_out线性投影,再reshape/permute/flatten还原为视频潜变量张量。
关键子模块逐一拆解
| 子模块 | 类 | 作用 | 源码位置 |
|---|---|---|---|
| 旋转位置编码 | WanRotaryPosEmbed | 对 T/H/W 三轴分别做 1D RoPE 并拼接 | L833 |
| Patch 嵌入 | nn.Conv3d× 2 | 视频潜变量与姿态潜变量分别 patchify 后相加 | L1087-L1088 |
| 条件嵌入 | WanTimeTextImageEmbedding | 时间步正弦嵌入 + 文本投影 + 图像嵌入 | L784 |
| 运动编码器 | WanAnimateMotionEncoder | 外观卷积网络 + 运动线性网络 + 线性运动分解(QR 正交化) | L282 |
| 面部编码器 | WanAnimateFaceEncoder | 时间因果 Conv1d + 多头(含 padding token)特征编码 | L356 |
| 主干块 | WanTransformerBlock | 自注意力 + 交叉注意力 + FFN,带 scale-shift 调制 | L899 |
| 面部适配器 | WanAnimateFaceBlockCrossAttention | 以运动特征为 KV 的临时对齐交叉注意力,QK 带 RMSNorm | L493 |
运动编码器内部(L282-L353)包含两条通路:外观编码器由MotionConv2d卷积残差块逐级下采样,最后接 4×4 卷积得到style_dim维外观特征;运动网络则由多层MotionLinear构成(源码注释特别提醒:这些线性层之间没有激活函数,这与官方实现保持一致)。最终的"线性运动分解"操作把运动向量做成对角矩阵,与motion_synthesis_weight的 QR 正交分解矩阵相乘再按行求和,得到运动向量。源码中 QR 分解被刻意 upcast 到 FP32 以保证数值稳定性(L340-L351)。
面部编码器(L386-L419)采用时间因果填充(time_causal_padding = (kernel_size - 1, 0),replicate模式)的 Conv1d 串接 LayerNorm + SiLU,最后一维拼接一个可学习的 padding token,把输出变成(B, T, N+1, C_out)——多出的第N+1个头用于吸收无运动帧。
面部适配器注入逻辑在 L1273-L1279:当block_idx % inject_face_latents_blocks == 0时取对应的face_adapter[block_idx // inject_face_latents_blocks]执行交叉注意力,并把输出与主干隐状态残差相加;考虑到模型并行场景,适配器输出会被显式搬到主干所在设备。
配置参数全解
以下参数均来自__init__的@register_to_config签名(transformer_wan_animate.py#L1040-L1069)与类 docstring(L992-L1023),Wan2.2-Animate-14B 官方权重即使用这些默认值:
| 参数 | 默认值 | 含义 |
|---|---|---|
patch_size | (1, 2, 2) | 视频嵌入的 3D patch 尺寸(t_patch, h_patch, w_patch) |
num_attention_heads | 40 | 注意力头数 |
attention_head_dim | 128 | 每个注意力头的通道数 |
in_channels | 36 | 输入通道数,等于2 * latent_channels + 4 |
latent_channels | 16 | Wan VAE 的潜变量通道数 |
out_channels | 16 | 输出通道数,缺省时回退为latent_channels |
text_dim | 4096 | 文本嵌入维度(对应 umT5 输出) |
freq_dim | 256 | 正弦时间嵌入维度 |
ffn_dim | 13824 | 前馈网络中间维度 |
num_layers | 40 | Transformer 块数量 |
cross_attn_norm | True | 是否启用交叉注意力归一化 |
qk_norm | "rms_norm_across_heads" | Q/K 归一化方式 |
eps | 1e-6 | 归一化层 epsilon |
image_dim | 1280 | 图像嵌入通道数(CLIP 视觉特征维度),None表示不投影 |
added_kv_proj_dim | None | 附加 KV 投影通道数,None表示不启用 |
rope_max_seq_len | 1024 | RoPE 最大序列长度 |
pos_embed_seq_len | None | 图像位置嵌入序列长度 |
motion_encoder_channel_sizes | 默认字典 | 运动编码器各分辨率通道数映射(512→32 至 4→512) |
motion_encoder_size | 512 | 面部视频输入分辨率(须为正方形) |
motion_style_dim | 512 | 外观特征维度 |
motion_dim | 20 | 运动向量维度 |
motion_encoder_dim | 512 | 运动编码器输出维度 |
face_encoder_hidden_dim | 1024 | 面部编码器隐藏维度 |
face_encoder_num_heads | 4 | 面部编码器头数 |
inject_face_latents_blocks | 5 | 每隔多少层注入一次面部特征 |
motion_encoder_batch_size | 8 | 面部视频按 batch 编码的批大小 |
其中in_channels与latent_channels二选一即可——构造时若只给其一,代码会自动按in_channels = 2 * latent_channels + 4推导(L1074-L1082),若两者都给则必须满足该恒等式,否则直接断言报错。
类级还声明了若干对下游生态至关重要的元属性(L1025-L1038):
_supports_gradient_checkpointing = True:支持训练时梯度检查点,blocks循环内会走_gradient_checkpointing_func分支;_no_split_modules = ["WanTransformerBlock", "MotionEncoderResBlock"]:accelerate 设备放置时按这些模块切分;_keep_in_fp32_modules:time_embedder、scale_shift_table、各norm、motion_synthesis_weight、rope等数值敏感模块保持在 FP32;_repeated_blocks = ["WanTransformerBlock"]:供循环权重共享与加载优化使用。
前向输入输出约定
forward的完整签名与形状约束(L1155-L1166):
| 参数 | 形状 | 说明 |
|---|---|---|
hidden_states | (B, 2C+4, T+1, H, W) | 含参考帧的加噪视频潜变量;C=16 为 Wan VAE 潜变量通道 |
timestep | (B,)的torch.LongTensor | 去噪循环当前时间步 |
encoder_hidden_states | (B, seq, 4096) | umT5 文本嵌入 |
encoder_hidden_states_image | (B, 257, 1280) | 参考人物图的 CLIP 视觉特征(可选) |
pose_hidden_states | (B, C, T, H, W) | 姿态视频潜变量,帧数须比hidden_states少 1,否则抛 ValueError(L1204-L1208) |
face_pixel_values | (B, 3, S, H', W') | 像素空间面部视频,S 为推理段帧数(官方为 77),H'/W' 须与motion_encoder_size一致(512×512) |
motion_encode_batch_size | int,可选 | 运动编码器的分批批大小,缺省用配置值 8 |
return_dict | bool,默认True | 返回字典还是元组 |
attention_kwargs | dict,可选 | 透传给 AttentionProcessor 的参数 |
输出为Transformer2DModelOutput(sample字段)或纯元组,sample为去噪后的视频潜变量。测试 test_models_transformer_wan_animate.py 中的真实尺寸 dummy 输入(hidden_states=(1,36,21,64,64)、pose_hidden_states=(1,16,20,64,64)、face_pixel_values=(1,3,77,512,512))完整印证了上述形状约定。
在 WanAnimatePipeline 中的实际调用
该模型通常不单独使用,而是作为 WanAnimatePipeline 的transformer组件被驱动(L28)。管道内置的示例(L50-L115)展示了两种模式:
import torch from diffusers import WanAnimatePipeline from diffusers.utils import export_to_video, load_image, load_video model_id = "Wan-AI/Wan2.2-Animate-14B-Diffusers" pipe = WanAnimatePipeline.from_pretrained(model_id, torch_dtype=torch.bfloat16) pipe.vae.to(torch.float32) # 可选:将 Wan VAE upcast 到 FP32 pipe.to("cuda") image = load_image("path/to/reference_character.png") pose_video = load_video("path/to/pose_video.mp4") face_video = load_video("path/to/face_video.mp4") # 动画模式:用姿态/面部视频驱动角色 output = pipe( image=image, pose_video=pose_video, face_video=face_video, prompt="A woman speaking naturally, high quality, ultrarealistic detail.", height=720, width=1280, segment_frame_length=77, # 每个推理段的数据帧长度 guidance_scale=1.0, # Wan Animate 一般不使用 CFG num_inference_steps=20, mode="animate", ).frames[0] export_to_video(output, "output_animation.mp4", fps=30) # 替换模式:在背景视频中替换角色(需要额外提供背景与掩码视频) output = pipe( image=image, pose_video=pose_video, face_video=face_video, background_video=load_video("path/to/background_video.mp4"), mask_video=load_video("path/to/mask_video.mp4"), # 黑色区域保留,白色区域生成 prompt="A woman speaking naturally.", height=720, width=1280, segment_frame_length=77, guidance_scale=1.0, num_inference_steps=20, mode="replace", ).frames[0]管道的去噪循环(L1113-L1140)揭示了模型的两个关键调用细节:
- 段式推理(segment-wise):长视频被切成
segment_frame_length=77帧的推理段,逐段去噪,段间通过前一帧潜变量做条件衔接(prev_segment_conditioning_frames),每一段都会调用一次self.transformer(...),并传入pose_hidden_states、face_pixel_values与motion_encode_batch_size; - CFG 的面部置空策略:若启用 classifier-free guidance,无条件分支会把面部视频像素整体置为
-1(face_pixel_values_uncond = face_video_segment * 0 - 1),其余输入保持不变——即"抹掉面部信息"而非"抹掉文本",这是面部驱动模型特有的 CFG 语义;同时cache_context("cond")/cache_context("uncond")的缓存上下文说明该模型支持 Diffusers 的条件缓存机制。
量化与单文件加载
从测试矩阵可以确认该模型对多种量化与编译生态的支持(见 test_models_transformer_wan_animate.py):
- BitsAndBytes 量化(
TestWanAnimateTransformer3DBitsAndBytes):支持 int8 量化推理,测试中真实维度输入为(1,36,5,16,16)等; - TorchAO 量化(
TestWanAnimateTransformer3DTorchAo):支持 TorchAO 低比特量化,使用 bf16 基座; - GGUF 量化(
TestWanAnimateTransformer3DGGUF/GGUFCompileTesterMixin):可通过单文件 GGUF 权重(如Wan2.2-Animate-14B-Q2_K.gguf)加载,并在 single_file_model.py#L172 中注册了WanAnimateTransformer3DModel的 single-file 映射; - torch.compile:
TorchCompileTesterMixin覆盖编译路径,但测试特意跳过F.pad(mode="replicate")(WanAnimateFaceEncoder内的时间因果填充)触发的 dynamo 图中断场景,说明该模块在编译时存在已知限制; - 内存与训练:
MemoryTesterMixin、TrainingTesterMixin(含梯度检查点验证)均有覆盖。
测试验证与延伸阅读
模型测试配置(WanAnimateTransformer3DTesterConfig,test_models_transformer_wan_animate.py#L39-L143)演示了如何用自定义小通道配置构建模型做快速验证,其中特意缩小motion_encoder_channel_sizes(否则运动编码器会占据测试模型绝大多数参数),并断言输出通道数少于输入(12→4),这正对应"主干只预测潜变量增量"的设计。
若需继续深入,可顺藤摸瓜阅读:
- Wan 管道 API 文档:
WanAnimatePipeline的完整参数说明与预处理要求(姿态关键点、面部特征需预先从参考视频中提取); - WanAnimate2Transformer3DModel API 文档:Wan2.2-Animate 的第二代变体,其实现位于 transformer_wan_animate_2.py,测试位于 test_models_transformer_wan_animate_2.py;
- models/init.py 与 transformers/init.py:查看该模型在
diffusers中的导出与模块注册方式。
总之,WanAnimateTransformer3DModel把"运动条件注入"做到了模型内部:运动编码器负责从人脸像素中提炼运动向量,面部编码器负责时间对齐,面部适配器负责把信号按固定间隔注入 40 层主干——理解这三者的协作方式,是自定义训练或二次开发 Wan Animate 类应用的前提。
【免费下载链接】diffusers🤗 Diffusers: State-of-the-art diffusion models for image, video, and audio generation in PyTorch.项目地址: https://gitcode.com/GitHub_Trending/di/diffusers
创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考