Diffusers documentation
SanaWMLTX2RefinerTransformer3DModel
SanaWMLTX2RefinerTransformer3DModel
作为 SanaWMPipeline 第二阶段使用的块因果自回归精化器 DiT，由 SanaWMLTX2Refiner 驱动。
其架构与 LTX2VideoTransformer3DModel 完全一致——相同的配置参数、相同的子模块、相同的参数名称——因此已发布的 LTX-2 检查点可以原样加载。前向传播有所不同：
- 仅运行视频流（跳过音频和音频/视频交叉注意力分支），
- self-attention 针对一个显式的滑动窗口 KV 缓存（
SanaWMRefinerKVCache）运行，该缓存保存了 attention sink 以及最近的细化历史，因此每个 block 的计算量是有界的，且总细化成本随视频长度线性增长， - 调用方提供视频 RoPE，这使得每个自回归窗口都能保留源视频中每一帧的绝对索引（参见 SanaWMLTX2RefinerTransformer3DModel.build_rotary_emb_for_absolute_positions()）。
该模型可通过以下方式加载：
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
类 diffusers.SanaWMLTX2RefinerTransformer3DModel
< source >( 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 )
参数
- in_channels (
int，默认为128) — 输入中的通道数。 - out_channels (
int，默认为128) — 输出中的通道数。 - patch_size (
int，默认为1) — 用于 patch 嵌入层的空间 patch 大小。 - patch_size_t (
int，默认为1) — 用于 patch 嵌入层的时间 patch 大小。 - num_attention_heads (
int，默认为32) — 多头注意力机制中使用的注意力头数量。 - attention_head_dim (
int，默认为128) — 每个头中的通道数。 - cross_attention_dim (
int，默认为4096) — 交叉注意力头的通道数。 - num_layers (
int，默认为48) — 要使用的 Transformer 块层数。 - activation_fn (
str，默认为"gelu-approximate") — 前馈网络中使用的激活函数。 - qk_norm (
str，默认为"rms_norm_across_heads") — 要使用的归一化层。 - rope_type (
str，默认为"interleaved") — 要使用的 RoPE 应用方式（"interleaved"或"split"）。
用作 SANA-WM 第二阶段的块因果自回归精化 Transformer。
在架构上与 LTX2VideoTransformer3DModel 完全一致——相同的配置参数、相同的子模块、相同的参数名称——因此已发布的 LTX-2 检查点可以原样加载到该模型中。两者的区别在于前向传播过程：
- 仅运行视频流（跳过音频和音频/视频交叉注意力分支），
- 自注意力机制针对一个显式的滑动窗口 KV 缓存（
SanaWMRefinerKVCache）运行，该缓存保存注意力 sink 以及最近的细化历史，因此每个 AR 块的细化成本是有界的，并随视频长度线性增长， - 调用方提供视频 RoPE，使每个 AR 窗口都能保留源视频中每一帧的绝对索引（参见 build_rotary_emb_for_absolute_positions()）。
其余参数与 LTX2VideoTransformer3DModel 一一对应。保留音频侧的参数和子模块，纯粹是为了让检查点中的音频权重能够往返；精化器的前向传播不会使用它们。
build_rotary_emb_for_absolute_positions
< source >( batch_size: intframe_positions: listheight: intwidth: intdevice: devicefps: float ) → tuple[torch.Tensor, torch.Tensor]
为显式指定的绝对潜在帧索引列表构建视频 RoPE。
SanaWMLTX2AudioVideoRotaryPosEmbed.prepare_video_coords 假定输入为连续的 torch.arange(num_frames)，这在双向推理中是可行的。而滑动窗口自回归精化器则需要保留每帧在源视频中的绝对索引，以便 RoPE 在 [sink + recent + active] 窗口中捕获正确的时间相位。
forward
< source >( 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
参数
- hidden_states (
torch.Tensor) — 当前活动块的补丁化视频潜变量，形状为(batch_size, num_video_tokens, in_channels)。 - encoder_hidden_states (
torch.Tensor) — 形状为(batch_size, text_seq_len, caption_channels)的文本嵌入。 - timestep (
torch.Tensor) — 形状为(batch_size, num_video_tokens)的时间步，已经过self.config.timestep_scale_multiplier缩放。 - video_rotary_emb (
tuple[torch.Tensor, torch.Tensor]) — 当前块绝对帧位置的(cos, sin)RoPE，由 build_rotary_emb_for_absolute_positions() 返回。 - encoder_attention_mask (
torch.Tensor, 可选) — 形状为(batch_size, text_seq_len)的乘法文本注意力掩码。 - kv_cache (
SanaWMRefinerKVCache，可选) — 滑动窗口 KV 缓存，保存各层的注意力 sink 以及最近的细化历史，外加共享的sink_peRoPE。只要设置了kv_cache_mode，就必须提供该参数。 - kv_cache_mode (
str, 可选) — 取值之一："inject"：关注[sink + history + current]的 K/V（去噪步骤）。"capture_pre_rope"：无前缀；将本次前向传播的 RoPE 之前的 K/V 记录到缓存中（仅使用一次，用于从原始的 stage-1 潜变量中初始化注意力汇聚点）。"inject_and_capture_post_rope"：关注[sink + history + current]的 K/V，并将该块的 post-RoPE K/V 记录到缓存中，以便追加到历史中。
当
None时，该块仅对当前 token 执行标准的完整自注意力。 - attention_kwargs (
dict, 可选) — 可选的关键字参数，将转发给 LoRA 缩放处理。 - return_dict (
bool, 可选，默认为True) — 是否返回 Transformer2DModelOutput 而不是普通元组。
返回值
Transformer2DModelOutput 或 tuple
活动块的预测速度，形状为 (batch_size, num_video_tokens, out_channels)。
单个 AR 块上仅针对视频的前向传播。
SanaWMRefinerKVCache
class diffusers.models.transformers.transformer_sana_wm_refiner.SanaWMRefinerKVCache
< source >( num_layers: int )
容器，每个 Transformer 块持有一个 SanaWMRefinerKVLayerCache，外加共享的 sink RoPE。
此实现遵循 SANA-WM 第二阶段精化器训练时所使用的 rf_shifted_sink KV 缓存契约。精化过程是块因果的：每次对 block_size 个潜在帧进行去噪，同时关注一个有界窗口 [attention sink + recent history + active block] 的 K/V。
sink_pe：用于 sink 帧的 (cos, sin) RoPE 元组。该元组在每个 AR 窗口内根据滑动的sink_rope_offset重新构建，使 sink 紧邻有界工作缓存之前。由于 RoPE 不依赖于层，因此在各层间共享。
SanaWMRefinerKVLayerCache
class diffusers.models.transformers.transformer_sana_wm_refiner.SanaWMRefinerKVLayerCache
< source >( )
SANA-WM 第二阶段块因果 AR 精化器的逐层 KV 缓存。
保存滑动窗口前缀的两个部分，供 refiner 的自注意力机制使用，此外还包含一个用于读取上一次前向传播所捕获的 K/V 的槽位。所有张量的形状均为 (batch_size, num_tokens, inner_dim)（即按头进行 unflatten 之前的布局），与 refiner 自注意力机制进行拼接时的布局相匹配。
sink_k_pre/sink_v：pre-RoPE 的注意力 sink 帧的 K/V，从原始 stage-1 潜变量中一次性捕获。它们以 pre-RoPE 形式存储，以便每个 AR 窗口都能在其自身偏移的 sink 偏移量上重新应用 RoPE（SanaWMRefinerKVCache.sink_pe）。history_k/history_v：post-RoPE 的 K/V，对应已经精化的近期帧，可直接按原样拼接。captured_k_pre/captured_v_pre和captured_k_post/captured_v_post：由捕获kv_cache_mode写入的回读槽位。
弹出由上一次前向传播捕获的 RoPE 之后的 K/V。
弹出由上一次前向传播捕获的预 RoPE K/V。
返回经过 RoPE 处理后的近期历史 K/V，若为空则返回 None。
返回 RoPE 之前的 sink K/V，如果尚未捕获（或为空），则返回 None。
存储当前前向传播中经过 RoPE 处理后的 K/V。
存储当前前向传播中生成的 RoPE 之前的 K/V。
存储经过 RoPE 处理后的近期历史 K/V。
存储预 RoPE 的 sink K/V。
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 的输出。