Diffusers documentation
SanaWMTransformer3DModel
SanaWMTransformer3DModel
一个 3D 扩散 Transformer（16 亿参数），用于相机控制下的图像到视频生成，作为 SanaWMPipeline 的第一阶段采样器。该 Transformer 结合了：
- 一个双向 GDN-Triton 线性注意力主干分支（深度 20，隐藏维度 2240，20 个注意力头），
- 一个 UCPE（Unified Camera Pose Embedding，统一相机位姿嵌入）相机控制分支，它接收所请求轨迹的 raymap + Plücker 表示，并且
- 一种 Wan 风格的 3D 旋转位置嵌入，外加每隔
softmax_every_n层注入一次的周期性 softmax 注意力块。
该模型可通过以下方式加载：
import torch
from diffusers import SanaWMTransformer3DModel
transformer = SanaWMTransformer3DModel.from_pretrained(
"Efficient-Large-Model/SANA-WM_bidirectional-diffusers",
subfolder="transformer",
torch_dtype=torch.bfloat16,
)SanaWMTransformer3DModel
class diffusers.SanaWMTransformer3DModel
< source >( in_channels: int = 128num_layers: int = 20hidden_size: int = 2240num_attention_heads: int = 20patch_size: tuple[int, int, int] = (1, 1, 1)softmax_every_n: int = 4linear_head_dim: int = 112ffn_type: str = 'GLUMBConvTemp't_kernel_size: int = 3conv_kernel_size: int = 4k_conv_only: bool = Truepos_embed_type: str = 'wan_rope'qk_norm: bool = Truecross_norm: bool = Truey_norm: bool = Truecam_attn_compress: int = 1init_cam_from_base: bool = Trueuse_chunk_plucker_post_attn: bool = Truechunk_plucker_channels: int = 48chunk_plucker_post_attn_blocks: int = 20fp32_attention: bool = Trueimage_size: int = 720caption_channels: int = 2304mlp_ratio: float = 3.0mlp_acts: tuple = ('silu', 'silu', None)use_pe: bool = Truelearn_sigma: bool = Falsepred_sigma: bool = Falsemixed_precision: str = 'bf16' )
参数
- in_channels (int，默认为 128) — VAE 潜在通道数（LTX-2）。
- softmax_every_n (int，默认为 4) — 每 N 个块使用一个 softmax 注意力块。
- linear_head_dim (int，默认为 112) — GDN 头维度。
- ffn_type（str，默认为
"GLUMBConvTemp"）— FFN。 - t_kernel_size (int，默认为 3) — 时间卷积核。
- conv_kernel_size (int，默认为 4) — 注意力机制中的空间卷积核大小。
- k_conv_only (bool，默认为 True) — 仅对 K 应用卷积。
- pos_embed_type (str，默认为
"wan_rope") — 位置嵌入。 - qk_norm (bool，默认为 True) — 对 Q/K 应用 RMSNorm。
- cross_norm (bool，默认为 True) — 对交叉注意力中的 K 应用 RMSNorm。
- y_norm（bool，默认为 True）— 对文本嵌入应用
attention_y_norm。 - init_cam_from_base (bool，默认为 True) — 未使用；相机分支从检查点加载。 保留此参数以便已发布的 config.json 文件能够正常加载。
- use_chunk_plucker_post_attn (bool，默认为 True)。 —
- chunk_plucker_channels（int，默认值为 48）—
6 dims * temporal_stride 8。 - chunk_plucker_post_attn_blocks (int，默认为 20) — 所有块。
- fp32_attention (bool，默认为 True) — 未使用；注意力始终在 fp32 下运行。保留此参数以便加载已发布的 config.json 文件。
- image_size (int，默认为 720) — 标称图像尺寸。
- caption_channels（int，默认为 2304）— Gemma-2 的隐藏层大小。
SANA-WM 1600M 双向相机控制 DiT。
单类 DiT（depth=20，hidden_size=2240，patch_size=(1,1,1)，num_heads=20，即公开的 Efficient-Large-Model/SANA-WM_bidirectional 发布版本）。save_pretrained / from_pretrained 通过 ConfigMixin 开箱即用。
每个块都会执行一个相机条件（UCPE）注意力机制：[BidirectionalGDNUCPESinglePathLiteLA]，但每隔 softmax_every_n 个块会执行其对应的 softmax 版本 [_SoftmaxUCPESinglePathLiteLA]。
state-dict 与公开的 sana 检查点完全相同，只是有意移除了 pos_embed 缓冲区。
forward
< source >( hidden_states: torch.Tensortimestep: torch.Tensorencoder_hidden_states: torch.Tensorencoder_attention_mask: torch.Tensor | None = Nonemask: torch.Tensor | None = Nonereturn_dict: bool = Truedata_info: Optional[dict] = Nonecamera_conditions: torch.Tensor | None = Nonechunk_plucker: torch.Tensor | None = Nonecam_pos_embeds: Optional[dict] = Nonepos_embeds: torch.Tensor | None = Noneraymats: torch.Tensor | None = Noneframe_valid_mask: torch.Tensor | None = Nonechunk_size: int | None = None )
参数
- hidden_states —
(B, C, T, H, W)潜在表示。 - timestep —
(B, 1, T)逐帧扩散时间步（LTX 风格）。 - encoder_hidden_states —
(B, 1, L, D_caption)文本嵌入。 - encoder_attention_mask —
(B, L)文本注意力掩码（diffusers 约定）。 - mask —
encoder_attention_mask的别名，用于匹配 sana DiT 的关键字参数名称。如果同时传入两者，mask优先。 - return_dict — 如果为
True（默认值），则返回Transformer2DModelOutput； 否则返回一个单元素元组(sample,)。 - data_info — 额外的条件信息；当
data_info["image_vae_embeds"]存在时，会沿通道轴与潜变量进行拼接。 - camera_conditions —
(B, T, 20)用于驱动相机控制（UCPE）分支的原始相机条件。 - chunk_plucker — Plucker 光线嵌入
(B, C, T, H, W)，当模型配置了use_chunk_plucker_input/use_chunk_plucker_post_attn时会被使用。 - cam_pos_embeds — 可选的预计算相机位置嵌入（
"absmap"/"P"条目），用于复用而非重新计算。 - pos_embeds — 可选的预计算旋转位置嵌入；当为
None时，将根据潜在张量的形状进行构建。 - raymats — 可选的预计算 UCPE 光线矩阵（仅在
cam_pos_embeds不包含"P"时使用）。 - frame_valid_mask — 可选的逐帧有效性掩码，形状为
(B, 1, T, 1, 1)、(B, 1, T)或(B, T)。 - chunk_size — 转发给各 block 的分块长度覆盖值；若未设置，则回退到每个 block 配置的
chunk_size。
运行 SANA-WM DiT。
x: (N, T, patch_size*2 C) imgs: (N, H, W, C)
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 的输出。