Hugging Face's logo

Diffusers documentation

SanaWMTransformer3DModel

Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

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_n layers.

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

< >

( 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_norm to 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

< >

( 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_mask matching the sana DiT’s kwarg name. If both are passed, mask takes precedence.
  • return_dict — If True (default), returns a Transformer2DModelOutput; 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 with use_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 None they are built from the latent shape.
  • raymats — Optional pre-computed UCPE ray matrices (used only when cam_pos_embeds does 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.

unpatchify

< >

( xpost_patch_num_framespost_patch_heightpost_patch_width )

x: (N, T, patch_size*2 C) imgs: (N, H, W, C)

Transformer2DModelOutput

class diffusers.models.modeling_outputs.Transformer2DModelOutput

< >

( sample: torch.Tensor )

Parameters

  • sample (torch.Tensor of 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 the encoder_hidden_states input. If discrete, returns probability distributions for the unnoised latent pixels.

The output of Transformer2DModel.

Update on GitHub

←SanaWMLTX2RefinerTransformer3DModel SD3Transformer2DModel→