Hugging Face's logo

Diffusers documentation

SanaWMLTX2RefinerTransformer3DModel

Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

SanaWMLTX2RefinerTransformer3DModel

The chunk-causal autoregressive refiner DiT used as stage 2 of SanaWMPipeline, driven by SanaWMLTX2Refiner.

It is architecturally identical to LTX2VideoTransformer3DModel — same config arguments, same submodules, same parameter names — so a released LTX-2 checkpoint loads into it unchanged. The forward pass differs:

  • only the video stream is run (the audio and audio/video cross-attention branches are skipped),
  • self-attention runs against an explicit sliding-window KV cache (SanaWMRefinerKVCache) holding the attention sink plus the recent refined history, so per-block compute is bounded and total refinement cost scales linearly with video length,
  • the caller supplies the video RoPE, which lets each autoregressive window keep every frame’s absolute index in the source video (see SanaWMLTX2RefinerTransformer3DModel.build_rotary_emb_for_absolute_positions()).

The model can be loaded with:

import torch
from diffusers import SanaWMLTX2RefinerTransformer3DModel

transformer = SanaWMLTX2RefinerTransformer3DModel.from_pretrained(
    "Efficient-Large-Model/SANA-WM_bidirectional-diffusers",
    subfolder="refiner/transformer",
    torch_dtype=torch.bfloat16,
)

SanaWMLTX2RefinerTransformer3DModel

class diffusers.SanaWMLTX2RefinerTransformer3DModel

< >

( in_channels: int = 128out_channels: int | None = 128patch_size: int = 1patch_size_t: int = 1num_attention_heads: int = 32attention_head_dim: int = 128cross_attention_dim: int = 4096vae_scale_factors: tuple = (8, 32, 32)pos_embed_max_pos: int = 20base_height: int = 2048base_width: int = 2048gated_attn: bool = Falsecross_attn_mod: bool = Falseaudio_in_channels: int = 128audio_out_channels: int | None = 128audio_patch_size: int = 1audio_patch_size_t: int = 1audio_num_attention_heads: int = 32audio_attention_head_dim: int = 64audio_cross_attention_dim: int = 2048audio_scale_factor: int = 4audio_pos_embed_max_pos: int = 20audio_sampling_rate: int = 16000audio_hop_length: int = 160audio_gated_attn: bool = Falseaudio_cross_attn_mod: bool = Falsenum_layers: int = 48activation_fn: str = 'gelu-approximate'qk_norm: str = 'rms_norm_across_heads'norm_elementwise_affine: bool = Falsenorm_eps: float = 1e-06caption_channels: int = 3840attention_bias: bool = Trueattention_out_bias: bool = Truerope_theta: float = 10000.0rope_double_precision: bool = Truecausal_offset: int = 1timestep_scale_multiplier: int = 1000cross_attn_timestep_scale_multiplier: int = 1000rope_type: str = 'interleaved'use_prompt_embeddings = Trueperturbed_attn: bool = Falseff_bias: bool = Trueaudio_ff_bias: bool = Trueuse_prompt_adaln_single: bool = Trueuse_keyframes_abs_pos_embedding: bool = False )

Parameters

  • in_channels (int, defaults to 128) — The number of channels in the input.
  • out_channels (int, defaults to 128) — The number of channels in the output.
  • patch_size (int, defaults to 1) — The size of the spatial patches to use in the patch embedding layer.
  • patch_size_t (int, defaults to 1) — The size of the temporal patches to use in the patch embedding layer.
  • num_attention_heads (int, defaults to 32) — The number of heads to use for multi-head attention.
  • attention_head_dim (int, defaults to 128) — The number of channels in each head.
  • cross_attention_dim (int, defaults to 4096) — The number of channels for cross attention heads.
  • num_layers (int, defaults to 48) — The number of layers of Transformer blocks to use.
  • activation_fn (str, defaults to "gelu-approximate") — Activation function to use in feed-forward.
  • qk_norm (str, defaults to "rms_norm_across_heads") — The normalization layer to use.
  • rope_type (str, defaults to "interleaved") — Which RoPE application to use ("interleaved" or "split").

The chunk-causal autoregressive refiner transformer used as SANA-WM stage 2.

Architecturally identical to LTX2VideoTransformer3DModel — same config arguments, same submodules, same parameter names — so a released LTX-2 checkpoint loads into it unchanged. What differs is the forward pass:

  • only the video stream is run (the audio and audio/video cross-attention branches are skipped),
  • self-attention runs against an explicit sliding-window KV cache (SanaWMRefinerKVCache) holding the attention sink plus recent refined history, so refinement cost is bounded per AR block and scales linearly with video length,
  • the caller supplies the video RoPE, which lets each AR window keep every frame’s absolute index in the source video (see build_rotary_emb_for_absolute_positions()).

The remaining arguments mirror LTX2VideoTransformer3DModel one-for-one. The audio-side arguments and submodules are kept purely so the checkpoint’s audio weights round-trip; they are not used by the refiner forward.

build_rotary_emb_for_absolute_positions

< >

( batch_size: intframe_positions: listheight: intwidth: intdevice: devicefps: float ) → tuple[torch.Tensor, torch.Tensor]

Parameters

  • batch_size (int) — Batch size to broadcast the coordinates to.
  • frame_positions (list[int]) — Absolute latent-frame indices covered by this window.
  • height (int), width (int) — Latent spatial resolution.
  • device (torch.device) — Device to build the coordinates on.
  • fps (float) — Video frame rate, which drives LTX-2’s temporal RoPE scaling.

Returns

tuple[torch.Tensor, torch.Tensor]

the (cos, sin) RoPE tuple.

Build the video RoPE for an explicit list of absolute latent-frame indices.

SanaWMLTX2AudioVideoRotaryPosEmbed.prepare_video_coords assumes a contiguous torch.arange(num_frames), which is fine for bidirectional inference. The sliding-window AR refiner instead needs to keep each frame’s absolute index in the source video, so RoPE captures the correct temporal phase across the [sink + recent + active] window.

forward

< >

( hidden_states: Tensorencoder_hidden_states: Tensortimestep: Tensorvideo_rotary_emb: tupleencoder_attention_mask: typing.Optional[torch.Tensor] = Nonekv_cache: diffusers.models.transformers.transformer_sana_wm_refiner.SanaWMRefinerKVCache | None = Nonekv_cache_mode: str | None = Noneattention_kwargs: dict | None = Nonereturn_dict: bool = True ) → Transformer2DModelOutput or tuple

Parameters

  • hidden_states (torch.Tensor) — Patchified video latents of the active block, of shape (batch_size, num_video_tokens, in_channels).
  • encoder_hidden_states (torch.Tensor) — Text embeddings of shape (batch_size, text_seq_len, caption_channels).
  • timestep (torch.Tensor) — Timestep of shape (batch_size, num_video_tokens), already scaled by self.config.timestep_scale_multiplier.
  • video_rotary_emb (tuple[torch.Tensor, torch.Tensor]) — The (cos, sin) RoPE for the active block’s absolute frame positions, as returned by build_rotary_emb_for_absolute_positions().
  • encoder_attention_mask (torch.Tensor, optional) — Multiplicative text attention mask of shape (batch_size, text_seq_len).
  • kv_cache (SanaWMRefinerKVCache, optional) — Sliding-window KV cache holding the per-layer attention sink and recent refined history, plus the shared sink_pe RoPE. Required whenever kv_cache_mode is set.
  • kv_cache_mode (str, optional) — One of:

    • "inject": attend to [sink + history + current] K/V (the denoising steps).
    • "capture_pre_rope": no prefix; record the pre-RoPE K/V of this forward into the cache (used once to seed the attention sink from the raw stage-1 latents).
    • "inject_and_capture_post_rope": attend to [sink + history + current] K/V and record this block’s post-RoPE K/V into the cache so it can be appended to the history.

    When None, the block runs plain full self-attention over the current tokens only.

  • attention_kwargs (dict, optional) — Optional kwargs forwarded to the LoRA scale handling.
  • return_dict (bool, optional, defaults to True) — Whether to return a Transformer2DModelOutput instead of a plain tuple.

Returns

Transformer2DModelOutput or tuple

the predicted velocity for the active block, of shape (batch_size, num_video_tokens, out_channels).

Video-only forward pass over a single AR block.

SanaWMRefinerKVCache

class diffusers.models.transformers.transformer_sana_wm_refiner.SanaWMRefinerKVCache

< >

( num_layers: int )

Parameters

  • num_layers (int) — Number of transformer blocks to allocate a per-layer cache for.

Container holding one [SanaWMRefinerKVLayerCache] per transformer block, plus the shared sink RoPE.

This implements the rf_shifted_sink KV-cache contract the SANA-WM stage-2 refiner was trained with. Refinement is chunk-causal: block_size latent frames are denoised at a time while attending to a bounded window of [attention sink + recent history + active block] K/V.

  • sink_pe: the (cos, sin) RoPE tuple for the sink frames, rebuilt per AR window at the sliding sink_rope_offset so the sink sits immediately before the bounded working cache. Shared across layers because RoPE does not depend on the layer.

SanaWMRefinerKVLayerCache

class diffusers.models.transformers.transformer_sana_wm_refiner.SanaWMRefinerKVLayerCache

< >

( )

Per-layer KV cache for the SANA-WM stage-2 chunk-causal AR refiner.

Holds the two halves of the sliding-window prefix that the refiner’s self-attention attends to, plus a slot for reading back the K/V that the last forward captured. All tensors are (batch_size, num_tokens, inner_dim) (i.e. before the per-head unflatten), matching the layout the refiner’s self-attention concatenates in.

  • sink_k_pre / sink_v: pre-RoPE K/V of the attention-sink frames, captured once from the raw stage-1 latents. They are stored pre-RoPE so each AR window can re-apply RoPE at its own shifted sink offset (SanaWMRefinerKVCache.sink_pe).
  • history_k / history_v: post-RoPE K/V of the already refined recent frames, ready to be concatenated as-is.
  • captured_k_pre / captured_v_pre and captured_k_post / captured_v_post: readback slots written by the capture kv_cache_modes.

get_captured_post_rope

< >

( )

Pop the post-RoPE K/V captured by the last forward.

get_captured_pre_rope

< >

( )

Pop the pre-RoPE K/V captured by the last forward.

get_history

< >

( )

Return the post-RoPE recent-history K/V, or None if empty.

get_sink

< >

( )

Return the pre-RoPE sink K/V, or None if it has not been captured (or is empty).

store_captured_post_rope

< >

( key: Tensorvalue: Tensor )

Store the post-RoPE K/V produced by the current forward.

store_captured_pre_rope

< >

( key: Tensorvalue: Tensor )

Store the pre-RoPE K/V produced by the current forward.

store_history

< >

( history_k: Tensorhistory_v: Tensor )

Store the post-RoPE recent-history K/V.

store_sink

< >

( sink_k_pre: Tensorsink_v: Tensor )

Store the pre-RoPE sink K/V.

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

←SanaVideoTransformer3DModel SanaWMTransformer3DModel→