Diffusers documentation
WanAnimateTransformer3DModel
WanAnimateTransformer3DModel
阿里巴巴 Wan 团队在 Wan Animate 中介绍了一种用于 3D 类视频数据的 Diffusion Transformer 模型。
可以使用以下代码片段加载该模型。
from diffusers import WanAnimateTransformer3DModel
transformer = WanAnimateTransformer3DModel.from_pretrained("Wan-AI/Wan2.2-Animate-14B-Diffusers", subfolder="transformer", dtype=torch.bfloat16)WanAnimateTransformer3DModel
class diffusers.WanAnimateTransformer3DModel
< source >( patch_size: tuple = (1, 2, 2)num_attention_heads: int = 40attention_head_dim: int = 128in_channels: int | None = 36latent_channels: int | None = 16out_channels: int | None = 16text_dim: int = 4096freq_dim: int = 256ffn_dim: int = 13824num_layers: int = 40cross_attn_norm: bool = Trueqk_norm: str | None = 'rms_norm_across_heads'eps: float = 1e-06image_dim: int | None = 1280added_kv_proj_dim: int | None = Nonerope_max_seq_len: int = 1024pos_embed_seq_len: int | None = Nonemotion_encoder_channel_sizes: dict[str, int] | None = Nonemotion_encoder_size: int = 512motion_style_dim: int = 512motion_dim: int = 20motion_encoder_dim: int = 512face_encoder_hidden_dim: int = 1024face_encoder_num_heads: int = 4inject_face_latents_blocks: int = 5motion_encoder_batch_size: int = 8 )
参数
- patch_size (
tuple[int]，默认为(1, 2, 2)) — 用于视频嵌入的 3D patch 维度 (t_patch, h_patch, w_patch)。 - num_attention_heads (
int，默认为40) — 文本嵌入的固定长度。 - attention_head_dim (
int，默认为128) — 每个头中的通道数。 - in_channels (
int，默认为16) — 输入数据的通道数。 - out_channels (
int，默认为16) — 输出中的通道数。 - text_dim (
int，默认为512) — 文本嵌入的输入维度。 - freq_dim (
int，默认为256) — 正弦时间嵌入的维度。 - ffn_dim (
int，默认为13824) — 前馈网络中的中间维度。 - num_layers (
int，默认为40) — 要使用的 transformer 块层数。 - window_size (
tuple[int]，默认为(-1, -1)) — 局部注意力的窗口大小（-1 表示全局注意力）。 - cross_attn_norm（
bool，默认为True）— 启用交叉注意力归一化。 - qk_norm (
bool，默认为True) — 启用查询/键归一化。 - eps (
float，默认为1e-6) — 归一化层的 epsilon 值。 - image_dim (
int，可选，默认为1280) — 用于图像嵌入的通道数。如果为None，则不使用投影。 - added_kv_proj_dim (
int，可选，默认为5120) — 用于附加的 key 和 value 投影的通道数。如果为None，则不使用投影。
用于 WanAnimate 模型中类视频数据的 Transformer 模型。
forward
< source >( hidden_states: Tensortimestep: LongTensorencoder_hidden_states: Tensorencoder_hidden_states_image: typing.Optional[torch.Tensor] = Nonepose_hidden_states: typing.Optional[torch.Tensor] = Noneface_pixel_values: typing.Optional[torch.Tensor] = Nonemotion_encode_batch_size: int | None = Nonereturn_dict: bool = Trueattention_kwargs: dict[str, typing.Any] | None = None ) → ~models.transformer_2d.Transformer2DModelOutput or tuple
参数
- hidden_states（形状为
(B, 2C + 4, T + 1, H, W)的torch.Tensor）— 形状为(B, 2C + 4, T + 1, H, W)的输入含噪视频潜变量，其中 B 为批大小，C 为潜变量通道数（Wan VAE 为 16），T 为推理片段中的潜变量帧数，H 为潜变量高度，W 为潜变量宽度。 - timestep — (
torch.LongTensor): 去噪循环中的当前时间步。 - encoder_hidden_states (
torch.Tensor) — 来自文本编码器的文本嵌入（Wan Animate 使用 umT5）。 - encoder_hidden_states_image (
torch.Tensor) — 参考（角色）图像的 CLIP 视觉特征。 - pose_hidden_states (
torch.Tensor，形状为(B, C, T, H, W)) — 姿态视频潜变量。TODO: 描述 - face_pixel_values（形状为
(B, C', S, H', W')的torch.Tensor）— 像素空间中的面部视频（非潜在空间）。通常 C’ = 3，H’ 和 W’ 为面部视频的像素高度/宽度。此处 S 为推理片段长度，通常设置为 77。 - motion_encode_batch_size (
int, 可选) — 通过运动编码器对面部视频进行批量编码时的批处理大小。如果未设置，则默认为self.config.motion_encoder_batch_size。 - return_dict (
bool，可选，默认为True) — 是否以字典或元组的形式返回输出。 - attention_kwargs (
dict, 可选) — 一个 kwargs 字典，如果指定，将传递给 diffusers.models.attention_processor 中定义的AttentionProcessor，该处理器在self.processor下定义。
返回值
~models.transformer_2d.Transformer2DModelOutput 或 tuple
如果 return_dict 为 True，则返回一个 ~models.transformer_2d.Transformer2DModelOutput，其 sample 为去噪后的视频潜在表示；否则，返回一个普通的 tuple，其第一个元素为该张量。
Wan2.2-Animate 变换器模型的前向传播。
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 的输出。