Diffusers documentation
SanaWMTransformer3DModel
SanaWMTransformer3DModel
A 3D Diffusion Transformer (1.6B parameters) for camera-controlled image-to-video generation, used as the stage-1 sampler of SanaWMPipeline. The transformer combines:
- a bidirectional GDN-Triton linear-attention main branch (depth 20, hidden 2240, 20 heads),
- a UCPE (Unified Camera Pose Embedding) camera-control branch that consumes a raymap + Plücker representation of the requested trajectory, and
- a Wan-style 3D rotary position embedding plus periodic softmax-attention blocks injected every
softmax_every_nlayers.
The model can be loaded with:
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' )
Parameters
- in_channels (int, defaults to 128) — VAE latent channels (LTX-2).
- softmax_every_n (int, defaults to 4) — Use a softmax attention block every N blocks.
- linear_head_dim (int, defaults to 112) — GDN head dimension.
- ffn_type (str, defaults to
"GLUMBConvTemp") — FFN. - t_kernel_size (int, defaults to 3) — Temporal conv kernel.
- conv_kernel_size (int, defaults to 4) — Spatial conv kernel inside attention.
- k_conv_only (bool, defaults to True) — Apply conv only on K.
- pos_embed_type (str, defaults to
"wan_rope") — Position embedding. - qk_norm (bool, defaults to True) — RMSNorm on Q/K.
- cross_norm (bool, defaults to True) — RMSNorm on cross-attention K.
- y_norm (bool, defaults to True) — Apply
attention_y_normto text embeddings. - init_cam_from_base (bool, defaults to True) — Unused; the camera branch is loaded from the checkpoint. Kept so released config.json files load.
- use_chunk_plucker_post_attn (bool, defaults to True). —
- chunk_plucker_channels (int, defaults to 48) —
6 dims * temporal_stride 8. - chunk_plucker_post_attn_blocks (int, defaults to 20) — All blocks.
- fp32_attention (bool, defaults to True) — Unused; attention always runs in fp32. Kept so released config.json files load.
- image_size (int, defaults to 720) — Nominal image size.
- caption_channels (int, defaults to 2304) — Gemma-2 hidden size.
SANA-WM 1600M bidirectional camera-controlled DiT.
A single-class DiT (depth=20, hidden_size=2240, patch_size=(1,1,1), num_heads=20 — i.e. the public Efficient-Large-Model/SANA-WM_bidirectional release). save_pretrained / from_pretrained work out of the
box via ConfigMixin.
Every block runs a camera-conditioned (UCPE) attention: [BidirectionalGDNUCPESinglePathLiteLA], except every softmax_every_n-th block which runs its softmax counterpart [_SoftmaxUCPESinglePathLiteLA].
The state-dict is identical to the public sana checkpoint apart from the intentionally-removed pos_embed buffer.
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 )
Parameters
- hidden_states —
(B, C, T, H, W)latents. - timestep —
(B, 1, T)per-frame diffusion timesteps (LTX style). - encoder_hidden_states —
(B, 1, L, D_caption)text embeddings. - encoder_attention_mask —
(B, L)text attention mask (diffusers convention). - mask — Alias for
encoder_attention_maskmatching the sana DiT’s kwarg name. If both are passed,masktakes precedence. - return_dict — If
True(default), returns aTransformer2DModelOutput; otherwise returns a one-tuple(sample,). - data_info — Extra conditioning;
data_info["image_vae_embeds"]is concatenated to the latents along the channel axis when present. - camera_conditions —
(B, T, 20)raw camera conditions driving the camera-control (UCPE) branch. - chunk_plucker — Plucker ray embeddings
(B, C, T, H, W), consumed when the model is configured withuse_chunk_plucker_input/use_chunk_plucker_post_attn. - cam_pos_embeds — Optional pre-computed camera positional embeddings
(
"absmap"/"P"entries) reused instead of recomputing them. - pos_embeds — Optional pre-computed rotary position embeddings; when
Nonethey are built from the latent shape. - raymats — Optional pre-computed UCPE ray matrices (used only when
cam_pos_embedsdoes not carry"P"). - frame_valid_mask — Optional per-frame validity mask, shaped
(B, 1, T, 1, 1),(B, 1, T)or(B, T). - chunk_size — Chunk length override forwarded to the blocks; falls back
to each block’s configured
chunk_size.
Run the SANA-WM DiT.
x: (N, T, patch_size*2 C) imgs: (N, H, W, C)
Transformer2DModelOutput
class diffusers.models.modeling_outputs.Transformer2DModelOutput
< source >( sample: torch.Tensor )
Parameters
- sample (
torch.Tensorof shape(batch_size, num_channels, height, width)or(batch size, num_vector_embeds - 1, num_latent_pixels)if Transformer2DModel is discrete) — The hidden states output conditioned on theencoder_hidden_statesinput. If discrete, returns probability distributions for the unnoised latent pixels.
The output of Transformer2DModel.