Diffusers documentation
WanAnimate2Transformer3DModel
WanAnimate2Transformer3DModel
用于 Wan-Animate-2（由阿里巴巴 Wan 团队开发）的 3D 类视频数据 Diffusion Transformer 模型。它通过上下文参考机制，利用驱动视频的动作来使角色图像动起来：每个片段首先执行一次参考路径（kv_cache_mode="extract"），以缓存每一层的参考 K/V；随后，去噪路径（kv_cache_mode="cached"）通过灵活的 BlockMask，同时对生成 token 和缓存的参考 token 进行注意力计算。
可以使用以下代码片段加载该模型。
from diffusers import WanAnimate2Transformer3DModel
transformer = WanAnimate2Transformer3DModel.from_pretrained("Wan-AI/Wan2.2-Animate-2-14B-Diffusers", subfolder="transformer", dtype=torch.bfloat16)WanAnimate2Transformer3DModel
class diffusers.WanAnimate2Transformer3DModel
< source >( patch_size: tuple = (1, 2, 2)text_len: int = 512in_dim: int = 36dim: int = 5120ffn_dim: int = 13824freq_dim: int = 256text_dim: int = 4096out_dim: int = 16num_heads: int = 40num_layers: int = 40cross_attn_norm: bool = Trueeps: float = 1e-06use_img_emb: bool = Truerefer_offset_t: int = 1refer_offset_h: int = 0refer_offset_w: int = -1refer_stride: int = 1 )
参数
- patch_size (tuple[int]，默认为 (1, 2, 2)) — 用于视频嵌入的 3D patch 维度 (t_patch, h_patch, w_patch)。
- text_len (int，默认为 512) — 文本嵌入的固定长度。
- in_dim（int，默认为 36） — 输入通道数（2 * latent_channels + 4，其中 4 为掩码通道）。
- dim (int，默认为 5120) — Transformer 中的通道数。
- ffn_dim (int，默认为 13824) — 前馈网络中的中间维度。
- freq_dim (int，默认为 256) — 正弦时间嵌入的维度。
- text_dim (int，默认为 4096) — 文本嵌入的输入维度。
- out_dim (int，默认为 16) — 输出中的通道数。
- num_heads (int，默认为 40) — 注意力头的数量。
- num_layers (int，默认为 40) — 要使用的 transformer 块层数。
- cross_attn_norm (bool，默认为 True) — 启用交叉注意力归一化。
- eps (float，默认为 1e-6) — 归一化层的 epsilon 值。
- use_img_emb (bool，默认为 True) — 是否使用 CLIP 图像嵌入。
- refer_offset_t (int，默认为 1) — 参考帧时间维度的 RoPE 偏移量。
- refer_offset_h (int，默认为 0) — 参考图像高度维度的 RoPE 偏移量。
- refer_offset_w (int，默认为 -1) — 参考图像宽度维度的 RoPE 偏移量。-1 表示使用生成网格的大小。
- refer_stride (int，默认为 1) — 在参考上应用 RoPE 时的步长。
用于 Wan-Animate-2 模型中类视频数据的 Transformer 模型。
Wan-Animate-2 采用基于 KV 缓存的上下文注意力机制：首先对参考视频进行编码（kv_cache_mode="extract"）以填充 [WanAnimate2KVCache]，随后在每个去噪步骤（kv_cache_mode="cached"）中，通过 flex BlockMask 同时对生成 token 和缓存的参考 K/V 进行联合注意力计算。因此，生成过程的自注意力仅在 flex 注意力后端上运行；模型中的其他所有注意力均可在任意后端上工作。
forward
< source >( hidden_states: listtimestep: Tensorencoder_hidden_states: listcondition_latents: listkv_cache: WanAnimate2KVCachekv_cache_mode: strseq_len: intencoder_hidden_states_image: typing.Optional[torch.Tensor] = Noneoffset_grid_sizes: typing.Optional[torch.Tensor] = Nonereference_grid_sizes: typing.Optional[torch.Tensor] = Noneorigin_len: int | None = Noneorigin_area: list[int] | None = Noneis_uncondtion: bool = Falsereturn_dict: bool = True ) → Transformer2DModelOutput or tuple(list[torch.Tensor])
参数
- hidden_states (
list[torch.Tensor]) — 本次传递的潜变量——当kv_cache_mode="extract"时为参考潜变量，当kv_cache_mode="cached"时为带噪生成潜变量。 - timestep (
torch.Tensor) — 去噪时间步。在kv_cache_mode="extract"模式下会被忽略，该模式使用固定的时间步 1。 - encoder_hidden_states (
list[torch.Tensor]) — 本次前向传播的文本嵌入。 - condition_latents (
list[torch.Tensor]) — 在 patch embedding 之前与hidden_states拼接的条件潜变量。 - kv_cache (
WanAnimate2KVCache) — 在kv_cache_mode="extract"下写入，在"cached"下读取。 - kv_cache_mode (
str) —"extract"执行参考路径并填充kv_cache；"cached"针对缓存的参考 token 执行去噪步骤。 - seq_len (
int) — 每个样本在经过 patch embedding 后必须保持的 token 数量。 - encoder_hidden_states_image (
torch.Tensor, 可选) — CLIP 图像嵌入，当模型配置了use_img_emb时使用。 - offset_grid_sizes (
torch.Tensor, 可选) — 参考潜变量的 Patch 网格，用于解析任何设置为 -1 的refer_offset_*。在kv_cache_mode="extract"下为必需项；在"cached"模式下，reference_grid_sizes描述相同的网格并用于替代。 - reference_grid_sizes (
torch.Tensor，可选) — 参考潜变量的补丁网格，用于参考旋转嵌入。在kv_cache_mode="cached"下为必填项。 - origin_len (
int, 可选) — 完整视频的帧数，上下文块掩码将基于此构建。在kv_cache_mode="cached"下为必填项。 - origin_area (
list[int], 可选) — 完整视频的空间尺寸[height, width]，上下文块掩码将基于此构建。 在kv_cache_mode="cached"下为必填项。 - is_uncondtion (
bool, 可选) — 是否是无分类器引导（classifier-free guidance）的无条件分支。 - return_dict (
bool，可选，默认为True) — 是否返回~models.transformer_2d.Transformer2DModelOutput而不是普通元组。
返回值
Transformer2DModelOutput 或 tuple(list[torch.Tensor])
每个输入潜变量对应的预测样本，未进行分块处理；如果 return_dict 为 False，则返回一个普通元组。
Transformer2DModelOutput
类 diffusers.models.modeling_outputs.Transformer2DModelOutput
< source >( sample: torch.Tensor )
参数
- sample（形状为
(batch_size, num_channels, height, width)的torch.Tensor，或者如果 Transformer2DModel 是离散的，则为(batch size, num_vector_embeds - 1, num_latent_pixels)）— 以encoder_hidden_states输入为条件的隐藏状态输出。如果是离散的，则返回未加噪潜在像素的概率分布。
Transformer2DModel 的输出。