# SanaWMLTX2RefinerTransformer3DModel

作为 [SanaWMPipeline](/docs/diffusers/main/zh/api/pipelines/sana_wm#diffusers.SanaWMPipeline) 第二阶段使用的块因果自回归精化器 DiT，由 [SanaWMLTX2Refiner](/docs/diffusers/main/zh/api/pipelines/sana_wm#diffusers.SanaWMLTX2Refiner) 驱动。

其架构与 [LTX2VideoTransformer3DModel](/docs/diffusers/main/zh/api/models/ltx2_video_transformer3d#diffusers.LTX2VideoTransformer3DModel) 完全一致——相同的配置参数、相同的子模块、相同的参数名称——因此已发布的 LTX-2 检查点可以原样加载。前向传播有所不同：

* 仅运行视频流（跳过音频和音频/视频交叉注意力分支），
* 自注意力机制针对一个显式的滑动窗口 KV 缓存（`SanaWMRefinerKVCache`）运行，该缓存保存注意力汇聚点以及最近的细化历史，因此每个块的计算量是有界的，且总细化成本随视频长度线性增长，
* 调用方提供视频 RoPE，这使得每个自回归窗口都能保留源视频中每一帧的绝对索引（参见 [SanaWMLTX2RefinerTransformer3DModel.build_rotary_emb_for_absolute_positions()](/docs/diffusers/main/zh/api/models/sana_wm_refiner_transformer3d#diffusers.SanaWMLTX2RefinerTransformer3DModel.build_rotary_emb_for_absolute_positions)）。

该模型可通过以下方式加载：

```python
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]]

#### diffusers.SanaWMLTX2RefinerTransformer3DModel[[diffusers.SanaWMLTX2RefinerTransformer3DModel]]

```python
diffusers.SanaWMLTX2RefinerTransformer3DModel(in_channels: int = 128, out_channels: int | None = 128, patch_size: int = 1, patch_size_t: int = 1, num_attention_heads: int = 32, attention_head_dim: int = 128, cross_attention_dim: int = 4096, vae_scale_factors: tuple = (8, 32, 32), pos_embed_max_pos: int = 20, base_height: int = 2048, base_width: int = 2048, gated_attn: bool = False, cross_attn_mod: bool = False, audio_in_channels: int = 128, audio_out_channels: int | None = 128, audio_patch_size: int = 1, audio_patch_size_t: int = 1, audio_num_attention_heads: int = 32, audio_attention_head_dim: int = 64, audio_cross_attention_dim: int = 2048, audio_scale_factor: int = 4, audio_pos_embed_max_pos: int = 20, audio_sampling_rate: int = 16000, audio_hop_length: int = 160, audio_gated_attn: bool = False, audio_cross_attn_mod: bool = False, num_layers: int = 48, activation_fn: str = 'gelu-approximate', qk_norm: str = 'rms_norm_across_heads', norm_elementwise_affine: bool = False, norm_eps: float = 1e-06, caption_channels: int = 3840, attention_bias: bool = True, attention_out_bias: bool = True, rope_theta: float = 10000.0, rope_double_precision: bool = True, causal_offset: int = 1, timestep_scale_multiplier: int = 1000, cross_attn_timestep_scale_multiplier: int = 1000, rope_type: str = 'interleaved', use_prompt_embeddings = True, perturbed_attn: bool = False, ff_bias: bool = True, audio_ff_bias: bool = True, use_prompt_adaln_single: bool = True, use_keyframes_abs_pos_embedding: bool = False)
```

[源代码](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_sana_wm_refiner.py#L1122)

**参数：**

in_channels (`int`，默认为 `128`)：输入中的通道数。

out_channels (`int`，默认为 `128`)：输出中的通道数。

patch_size (`int`，默认为 `1`)：在 patch embedding 层中使用的空间 patch 的大小。

patch_size_t (`int`，默认为 `1`)：在 patch embedding 层中使用的时序 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](/docs/diffusers/main/zh/api/models/ltx2_video_transformer3d#diffusers.LTX2VideoTransformer3DModel) 完全一致——相同的配置参数、相同的子模块、相同的参数名称——因此已发布的 LTX-2 检查点可以原样加载到该模型中。两者的区别在于前向传播过程：

* 仅运行视频流（跳过音频和音频/视频交叉注意力分支），
* 自注意力机制针对一个显式的滑动窗口 KV 缓存（`SanaWMRefinerKVCache`）运行，该缓存包含注意力汇聚点（attention sink）以及最近的细化历史，因此每个 AR 块的细化成本是有界的，并且随视频长度线性增长，
* 调用方提供视频 RoPE，这使得每个 AR 窗口都能保留源视频中每一帧的绝对索引（参见 [build_rotary_emb_for_absolute_positions()](/docs/diffusers/main/zh/api/models/sana_wm_refiner_transformer3d#diffusers.SanaWMLTX2RefinerTransformer3DModel.build_rotary_emb_for_absolute_positions)）。

其余参数与 [LTX2VideoTransformer3DModel](/docs/diffusers/main/zh/api/models/ltx2_video_transformer3d#diffusers.LTX2VideoTransformer3DModel) 一一对应。保留音频侧的参数和子模块纯粹是为了让检查点中的音频权重能够往返；refiner 的前向传播不会使用它们。

#### build_rotary_emb_for_absolute_positions[[diffusers.SanaWMLTX2RefinerTransformer3DModel.build_rotary_emb_for_absolute_positions]]

```python
build_rotary_emb_for_absolute_positions(batch_size: int, frame_positions: list, height: int, width: int, device: device, fps: float)
```

[源代码](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_sana_wm_refiner.py#L1388)

**参数：**

batch_size (`int`) : 用于广播坐标的批次大小。

frame_positions (`list[int]`) : 此窗口所覆盖的绝对潜在帧索引。

height (`int`)，宽度 (`int`)：潜在空间分辨率。

device (`torch.device`) : 用于构建坐标的设备。

fps (`float`)：视频帧率，用于驱动 LTX-2 的时间 RoPE 缩放。

**返回：** `tuple[torch.Tensor, torch.Tensor]`

`(cos, sin)` RoPE 元组。

为显式指定的绝对潜在帧索引列表构建视频 RoPE。

`SanaWMLTX2AudioVideoRotaryPosEmbed.prepare_video_coords` 假定输入是连续的 `torch.arange(num_frames)`，这在双向推理中没有问题。而滑动窗口自回归精化器需要保留每帧在源视频中的绝对索引，以便 RoPE 在 `[sink + recent + active]` 窗口中捕获正确的时间相位。

#### forward[[diffusers.SanaWMLTX2RefinerTransformer3DModel.forward]]

```python
forward(hidden_states: Tensor, encoder_hidden_states: Tensor, timestep: Tensor, video_rotary_emb: tuple, encoder_attention_mask: typing.Optional[torch.Tensor] = None, kv_cache: diffusers.models.transformers.transformer_sana_wm_refiner.SanaWMRefinerKVCache | None = None, kv_cache_mode: str | None = None, attention_kwargs: dict | None = None, return_dict: bool = True)
```

[源代码](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_sana_wm_refiner.py#L1445)

**参数：**

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]`) : 由 [build_rotary_emb_for_absolute_positions()](/docs/diffusers/main/zh/api/models/sana_wm_refiner_transformer3d#diffusers.SanaWMLTX2RefinerTransformer3DModel.build_rotary_emb_for_absolute_positions) 返回的，用于活动块绝对帧位置的 `(cos, sin)` RoPE。

encoder_attention_mask (`torch.Tensor`, *可选*) : 形状为 `(batch_size, text_seq_len)` 的乘法文本注意力掩码。

kv_cache (`SanaWMRefinerKVCache`, *可选*) : 滑动窗口 KV 缓存，保存每层的注意力汇聚点（attention sink）和最近的细化历史，以及共享的 `sink_pe` RoPE。只要设置了 `kv_cache_mode`，就必须提供此项。

kv_cache_mode (`str`, *可选*) : 以下之一： - `"inject"`：关注 `[sink + history + current]` K/V（去噪步骤）。 - `"capture_pre_rope"`：无前缀；将本次前向传播的 pre-RoPE K/V 记录到缓存中（仅用于从原始 stage-1 潜变量初始化注意力 sink）。 - `"inject_and_capture_post_rope"`：关注 `[sink + history + current]` K/V，并将该块的 post-RoPE K/V 记录到缓存中，以便追加到 history。 当为 `None` 时，该块仅对当前 token 执行标准的完整自注意力。

attention_kwargs (`dict`, *可选*) : 转发给 LoRA 缩放处理的可选 kwargs。

return_dict (`bool`，*可选*，默认为 `True`)：是否返回 [Transformer2DModelOutput](/docs/diffusers/main/zh/api/models/cogview3plus_transformer2d#diffusers.models.modeling_outputs.Transformer2DModelOutput) 而非普通元组。

**返回：** [Transformer2DModelOutput](/docs/diffusers/main/zh/api/models/cogview3plus_transformer2d#diffusers.models.modeling_outputs.Transformer2DModelOutput) 或 `tuple`

活动块的预测速度，形状为 `(batch_size, num_video_tokens, out_channels)`。

单个 AR 块上仅针对视频的前向传播。

## SanaWMRefinerKVCache[[diffusers.models.transformers.transformer_sana_wm_refiner.SanaWMRefinerKVCache]]

#### diffusers.models.transformers.transformer_sana_wm_refiner.SanaWMRefinerKVCache[[diffusers.models.transformers.transformer_sana_wm_refiner.SanaWMRefinerKVCache]]

```python
diffusers.models.transformers.transformer_sana_wm_refiner.SanaWMRefinerKVCache(num_layers: int)
```

[源代码](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_sana_wm_refiner.py#L757)

**参数：**

num_layers (*int*) : 需要为每个 Transformer 块分配逐层缓存的块数。

容器，每个 Transformer 块持有一个 *SanaWMRefinerKVLayerCache*，外加共享的 sink RoPE。

此实现遵循 SANA-WM 第二阶段精化器训练时所使用的 `rf_shifted_sink` KV 缓存契约。精化过程是块因果的：每次对 *block_size* 个潜在帧进行去噪，同时关注一个有界窗口 *[注意力汇聚点 + 近期历史 + 活动块]* 的 K/V。

- `sink_pe`：用于 sink 帧的 *(cos, sin)* RoPE 元组。它在每个 AR 窗口中根据滑动的 `sink_rope_offset` 重新构建，使 sink 紧邻有界工作缓存之前。由于 RoPE 不依赖于层，因此在各层间共享。

## SanaWMRefinerKVLayerCache[[diffusers.models.transformers.transformer_sana_wm_refiner.SanaWMRefinerKVLayerCache]]

#### diffusers.models.transformers.transformer_sana_wm_refiner.SanaWMRefinerKVLayerCache[[diffusers.models.transformers.transformer_sana_wm_refiner.SanaWMRefinerKVLayerCache]]

```python
diffusers.models.transformers.transformer_sana_wm_refiner.SanaWMRefinerKVLayerCache()
```

[源代码](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_sana_wm_refiner.py#L668)

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` 写入的回读槽位。

#### get_captured_post_rope[[diffusers.models.transformers.transformer_sana_wm_refiner.SanaWMRefinerKVLayerCache.get_captured_post_rope]]

```python
get_captured_post_rope()
```

[源代码](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_sana_wm_refiner.py#L736)

弹出由上一次前向传播捕获的 RoPE 之后的 K/V。

#### get_captured_pre_rope[[diffusers.models.transformers.transformer_sana_wm_refiner.SanaWMRefinerKVLayerCache.get_captured_pre_rope]]

```python
get_captured_pre_rope()
```

[源代码](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_sana_wm_refiner.py#L722)

弹出由上一次前向传播捕获的预 RoPE K/V。

#### get_history[[diffusers.models.transformers.transformer_sana_wm_refiner.SanaWMRefinerKVLayerCache.get_history]]

```python
get_history()
```

[源代码](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_sana_wm_refiner.py#L711)

返回经过 RoPE 处理后的近期历史 K/V，若为空则返回 `None`。

#### get_sink[[diffusers.models.transformers.transformer_sana_wm_refiner.SanaWMRefinerKVLayerCache.get_sink]]

```python
get_sink()
```

[源代码](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_sana_wm_refiner.py#L700)

返回 RoPE 之前的 sink K/V，如果尚未捕获（或为空），则返回 `None`。

#### store_captured_post_rope[[diffusers.models.transformers.transformer_sana_wm_refiner.SanaWMRefinerKVLayerCache.store_captured_post_rope]]

```python
store_captured_post_rope(key: Tensor, value: Tensor)
```

[源代码](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_sana_wm_refiner.py#L731)

存储当前前向传播中经过 RoPE 处理后的 K/V。

#### store_captured_pre_rope[[diffusers.models.transformers.transformer_sana_wm_refiner.SanaWMRefinerKVLayerCache.store_captured_pre_rope]]

```python
store_captured_pre_rope(key: Tensor, value: Tensor)
```

[源代码](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_sana_wm_refiner.py#L717)

存储当前前向传播中生成的 RoPE 之前的 K/V。

#### store_history[[diffusers.models.transformers.transformer_sana_wm_refiner.SanaWMRefinerKVLayerCache.store_history]]

```python
store_history(history_k: Tensor, history_v: Tensor)
```

[源代码](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_sana_wm_refiner.py#L706)

存储经过 RoPE 处理后的近期历史 K/V。

#### store_sink[[diffusers.models.transformers.transformer_sana_wm_refiner.SanaWMRefinerKVLayerCache.store_sink]]

```python
store_sink(sink_k_pre: Tensor, sink_v: Tensor)
```

[源代码](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_sana_wm_refiner.py#L695)

存储预 RoPE 的 sink K/V。

## Transformer2DModelOutput[[diffusers.models.modeling_outputs.Transformer2DModelOutput]]

#### diffusers.models.modeling_outputs.Transformer2DModelOutput[[diffusers.models.modeling_outputs.Transformer2DModelOutput]]

```python
diffusers.models.modeling_outputs.Transformer2DModelOutput(sample: torch.Tensor)
```

[源代码](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/modeling_outputs.py#L21)

**参数：**

sample (`torch.Tensor`，形状为 `(batch_size, num_channels, height, width)`；若 [Transformer2DModel](/docs/diffusers/main/zh/api/models/transformer2d#diffusers.Transformer2DModel) 为离散模式，则为 `(batch size, num_vector_embeds - 1, num_latent_pixels)`)：以 `encoder_hidden_states` 输入为条件的隐藏状态输出。若为离散模式，则返回未加噪潜在像素的概率分布。

[Transformer2DModel](/docs/diffusers/main/zh/api/models/transformer2d#diffusers.Transformer2DModel) 的输出。