# SanaWMLTX2RefinerTransformer3DModel

DiT tự hồi quy nhân quả theo chunk (chunk-causal autoregressive refiner) được sử dụng làm giai đoạn 2 của [SanaWMPipeline](/docs/diffusers/main/vi/api/pipelines/sana_wm#diffusers.SanaWMPipeline), được điều khiển bởi [SanaWMLTX2Refiner](/docs/diffusers/main/vi/api/pipelines/sana_wm#diffusers.SanaWMLTX2Refiner).

Về mặt kiến trúc, nó giống hệt [LTX2VideoTransformer3DModel](/docs/diffusers/main/vi/api/models/ltx2_video_transformer3d#diffusers.LTX2VideoTransformer3DModel) — cùng các đối số cấu hình, cùng các module con, cùng tên tham số — do đó một checkpoint LTX-2 đã phát hành có thể được nạp vào mà không cần thay đổi. Quá trình truyền tiến (forward pass) khác nhau:

* chỉ luồng video được chạy (các nhánh chú ý chéo giữa âm thanh và âm thanh/video bị bỏ qua),
* self-attention chạy trên một KV cache trượt cửa sổ tường minh (`SanaWMRefinerKVCache`) chứa attention sink cùng lịch sử tinh chỉnh gần nhất, do đó chi phí tính toán cho mỗi khối bị giới hạn và tổng chi phí tinh chỉnh tăng tuyến tính theo độ dài video,
* người gọi cung cấp video RoPE, cho phép mỗi cửa sổ tự hồi quy giữ lại chỉ số tuyệt đối của từng khung hình trong video nguồn (xem [SanaWMLTX2RefinerTransformer3DModel.build_rotary_emb_for_absolute_positions()](/docs/diffusers/main/vi/api/models/sana_wm_refiner_transformer3d#diffusers.SanaWMLTX2RefinerTransformer3DModel.build_rotary_emb_for_absolute_positions)).

Mô hình có thể được tải bằng cách:

```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[int, int, int] = (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)
```

[Nguồn](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_sana_wm_refiner.py#L1110)

**Tham số:**

in_channels (`int`, mặc định là `128`) : Số lượng kênh trong đầu vào.

out_channels (`int`, mặc định là `128`) : Số lượng kênh trong đầu ra.

patch_size (`int`, mặc định là `1`) : Kích thước của các patch không gian được sử dụng trong lớp nhúng patch.

patch_size_t (`int`, mặc định là `1`) : Kích thước của các patch thời gian được sử dụng trong lớp nhúng patch.

num_attention_heads (`int`, mặc định là `32`) : Số lượng đầu (heads) được sử dụng cho chú ý đa đầu (multi-head attention).

attention_head_dim (`int`, mặc định là `128`) : Số lượng kênh trong mỗi head.

cross_attention_dim (`int`, mặc định là `4096`) : Số kênh cho các đầu chú ý chéo (cross attention).

num_layers (`int`, mặc định là `48`) : Số lượng lớp khối Transformer cần sử dụng.

activation_fn (`str`, mặc định là `"gelu-approximate"`) : Hàm kích hoạt được sử dụng trong mạng truyền thẳng.

qk_norm (`str`, mặc định là `"rms_norm_across_heads"`) : Lớp chuẩn hóa cần sử dụng.

rope_type (`str`, mặc định là `"interleaved"`) : Phương thức áp dụng RoPE cần sử dụng (`"interleaved"` hoặc `"split"`).

Transformer tinh chỉnh tự hồi quy theo chunk-causal được sử dụng làm giai đoạn 2 của SANA-WM.

Về mặt kiến trúc, hoàn toàn giống với [LTX2VideoTransformer3DModel](/docs/diffusers/main/vi/api/models/ltx2_video_transformer3d#diffusers.LTX2VideoTransformer3DModel) — cùng các đối số cấu hình, cùng các module con, cùng tên tham số — do đó một checkpoint LTX-2 đã phát hành có thể được nạp vào mà không cần thay đổi. Điều khác biệt nằm ở quá trình truyền tiến (forward pass):

* chỉ luồng video được chạy (các nhánh chú ý chéo giữa âm thanh và âm thanh/video bị bỏ qua),
* self-attention chạy trên một bộ nhớ đệm KV dạng cửa sổ trượt rõ ràng (`SanaWMRefinerKVCache`) chứa attention sink cùng lịch sử tinh chỉnh gần nhất, do đó chi phí tinh chỉnh bị giới hạn cho mỗi khối AR và tăng tuyến tính theo độ dài video,
* người gọi cung cấp video RoPE, cho phép mỗi cửa sổ AR giữ lại chỉ số tuyệt đối của từng khung hình trong video nguồn (xem [build_rotary_emb_for_absolute_positions()](/docs/diffusers/main/vi/api/models/sana_wm_refiner_transformer3d#diffusers.SanaWMLTX2RefinerTransformer3DModel.build_rotary_emb_for_absolute_positions)).

Các đối số còn lại phản ánh [LTX2VideoTransformer3DModel](/docs/diffusers/main/vi/api/models/ltx2_video_transformer3d#diffusers.LTX2VideoTransformer3DModel) một cách tương ứng. Các đối số và mô-đun con phía âm thanh được giữ lại hoàn toàn để trọng số âm thanh của checkpoint có thể được lưu trữ và tải lại; chúng không được sử dụng bởi tiến trình forward của 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[int], height: int, width: int, device: device, fps: float)
```

[Nguồn](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_sana_wm_refiner.py#L1376)

**Tham số:**

batch_size (`int`) : Kích thước batch để phát sóng các tọa độ.

frame_positions (`list[int]`) : Các chỉ số khung tiềm năng tuyệt đối được bao phủ bởi cửa sổ này.

height (`int`), width (`int`) : Độ phân giải không gian của latent.

device (`torch.device`) : Thiết bị để xây dựng các tọa độ trên đó.

fps (`float`) : Tần số khung hình của video, quyết định việc mở rộng RoPE theo thời gian của LTX-2.

**Trả về:** `tuple[torch.Tensor, torch.Tensor]`

bộ đôi RoPE `(cos, sin)`.

Xây dựng video RoPE cho một danh sách rõ ràng các chỉ số khung tiềm năng tuyệt đối.

`SanaWMLTX2AudioVideoRotaryPosEmbed.prepare_video_coords` giả định một `torch.arange(num_frames)` liên tục, điều này phù hợp cho suy luận hai chiều. Bộ tinh chỉnh AR cửa sổ trượt thay vào đó cần giữ chỉ số tuyệt đối của từng khung hình trong video nguồn, để RoPE nắm bắt đúng pha thời gian xuyên suốt cửa sổ `[sink + recent + active]`.

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

```python
forward(hidden_states: Tensor, encoder_hidden_states: Tensor, timestep: Tensor, video_rotary_emb: tuple[torch.Tensor, torch.Tensor], 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)
```

[Nguồn](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_sana_wm_refiner.py#L1433)

**Tham số:**

hidden_states (`torch.Tensor`) : Các latents video đã được patchify của khối đang hoạt động, có hình dạng `(batch_size, num_video_tokens, in_channels)`.

encoder_hidden_states (`torch.Tensor`) : Nhúng văn bản có hình dạng `(batch_size, text_seq_len, caption_channels)`.

timestep (`torch.Tensor`) : Timestep có shape `(batch_size, num_video_tokens)`, đã được nhân với `self.config.timestep_scale_multiplier`.

video_rotary_emb (`tuple[torch.Tensor, torch.Tensor]`) : RoPE `(cos, sin)` cho các vị trí tuyệt đối của khung tham chiếu khối đang hoạt động, như được trả về bởi [build_rotary_emb_for_absolute_positions()](/docs/diffusers/main/vi/api/models/sana_wm_refiner_transformer3d#diffusers.SanaWMLTX2RefinerTransformer3DModel.build_rotary_emb_for_absolute_positions).

encoder_attention_mask (`torch.Tensor`, *tùy chọn*) : Mặt nạ chú ý văn bản nhân có hình dạng `(batch_size, text_seq_len)`.

kv_cache (`SanaWMRefinerKVCache`, *tùy chọn*) : Bộ đệm KV dạng cửa sổ trượt, lưu trữ attention sink và lịch sử tinh chỉnh gần nhất cho từng lớp, cùng với RoPE `sink_pe` dùng chung. Bắt buộc khi `kv_cache_mode` được đặt.

kv_cache_mode (`str`, *tùy chọn*) : Một trong các giá trị: - `"inject"`: chú ý đến K/V của `[sink + history + current]` (các bước khử nhiễu). - `"capture_pre_rope"`: không có tiền tố; ghi lại K/V trước RoPE của bước forward này vào bộ nhớ đệm (được dùng một lần để khởi tạo attention sink từ các latent stage-1 thô). - `"inject_and_capture_post_rope"`: chú ý đến K/V của `[sink + history + current]` và ghi lại K/V sau RoPE của khối này vào bộ nhớ đệm để có thể thêm vào lịch sử. Khi `None`, khối chạy self-attention đầy đủ thông thường chỉ trên các token hiện tại.

attention_kwargs (`dict`, *tùy chọn*) : Các kwargs tùy chọn được chuyển tiếp đến việc xử lý tỷ lệ LoRA.

return_dict (`bool`, *tùy chọn*, mặc định là `True`) : Có nên trả về một [Transformer2DModelOutput](/docs/diffusers/main/vi/api/models/cogview3plus_transformer2d#diffusers.models.modeling_outputs.Transformer2DModelOutput) thay vì một tuple thông thường hay không.

**Trả về:** [Transformer2DModelOutput](/docs/diffusers/main/vi/api/models/cogview3plus_transformer2d#diffusers.models.modeling_outputs.Transformer2DModelOutput) hoặc `tuple`

tốc độ dự đoán cho khối hoạt động, có hình dạng `(batch_size, num_video_tokens, out_channels)`.

Đi qua tiến trình forward chỉ với video trên một khối AR đơn lẻ.

## 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)
```

[Nguồn](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_sana_wm_refiner.py#L745)

**Tham số:**

num_layers (*int*) : Số lượng khối transformer để cấp phát bộ nhớ đệm cho từng lớp.

Container chứa một [*SanaWMRefinerKVLayerCache*] cho mỗi khối transformer, cùng với RoPE sink dùng chung.

Phần này triển khai hợp đồng KV-cache `rf_shifted_sink` mà bộ tinh chỉnh giai đoạn 2 của SANA-WM được huấn luyện với. Quá trình tinh chỉnh là nhân quả theo khối: *block_size* khung latent được khử nhiễu mỗi lần trong khi chú ý đến một cửa sổ có giới hạn gồm *[attention sink + lịch sử gần đây + khối đang hoạt động]* K/V.

- `sink_pe`: bộ đôi *(cos, sin)* RoPE cho các khung sink, được xây dựng lại cho mỗi cửa sổ AR tại `sink_rope_offset` trượt để khung sink nằm ngay trước bộ nhớ đệm làm việc có giới hạn. Được dùng chung giữa các lớp vì RoPE không phụ thuộc vào lớp.

## 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()
```

[Nguồn](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_sana_wm_refiner.py#L656)

Bộ nhớ đệm KV theo từng lớp cho bộ tinh chỉnh AR nhân quả theo khối ở giai đoạn 2 của SANA-WM.

Chứa hai phần của tiền tố cửa sổ trượt mà self-attention của refiner chú ý đến, cùng một slot để đọc lại K/V mà lần forward cuối cùng đã ghi nhận. Tất cả các tensor đều có dạng *(batch_size, num_tokens, inner_dim)* (tức là trước khi unflatten theo từng head), khớp với bố cục mà self-attention của refiner nối vào.

- `sink_k_pre` / `sink_v`: **pre**-RoPE K/V của các khung attention-sink, được ghi nhận một lần từ các latents thô của giai đoạn 1. Chúng được lưu ở dạng pre-RoPE để mỗi cửa sổ AR có thể áp dụng lại RoPE tại độ lệch sink đã dịch chuyển của riêng nó (`SanaWMRefinerKVCache.sink_pe`).
- `history_k` / `history_v`: K/V **post**-RoPE của các khung gần đây đã được tinh chỉnh, sẵn sàng để nối trực tiếp.
- `captured_k_pre` / `captured_v_pre` và `captured_k_post` / `captured_v_post`: các ô đọc lại được ghi bởi các `kv_cache_mode` bắt.

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

```python
get_captured_post_rope()
```

[Nguồn](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_sana_wm_refiner.py#L724)

Lấy ra K/V sau RoPE được ghi nhận bởi bước forward cuối cùng.

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

```python
get_captured_pre_rope()
```

[Nguồn](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_sana_wm_refiner.py#L710)

Lấy ra K/V trước RoPE được ghi nhận bởi bước forward cuối cùng.

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

```python
get_history()
```

[Nguồn](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_sana_wm_refiner.py#L699)

Trả về K/V lịch sử gần đây sau RoPE, hoặc `None` nếu trống.

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

```python
get_sink()
```

[Nguồn](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_sana_wm_refiner.py#L688)

Trả về K/V sink trước RoPE, hoặc `None` nếu nó chưa được ghi nhận (hoặc rỗng).

#### 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)
```

[Nguồn](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_sana_wm_refiner.py#L719)

Lưu trữ K/V sau RoPE được tạo ra bởi bước forward hiện tại.

#### 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)
```

[Nguồn](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_sana_wm_refiner.py#L705)

Lưu trữ K/V trước RoPE được tạo ra bởi bước forward hiện tại.

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

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

[Nguồn](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_sana_wm_refiner.py#L694)

Lưu trữ K/V lịch sử gần đây sau RoPE.

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

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

[Nguồn](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_sana_wm_refiner.py#L683)

Lưu trữ K/V sink trước RoPE.

## 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)
```

[Nguồn](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/modeling_outputs.py#L21)

**Tham số:**

sample (`torch.Tensor` có shape `(batch_size, num_channels, height, width)` hoặc `(batch size, num_vector_embeds - 1, num_latent_pixels)` nếu [Transformer2DModel](/docs/diffusers/main/vi/api/models/transformer2d#diffusers.Transformer2DModel) là discrete) : Đầu ra của hidden states được điều kiện bởi đầu vào `encoder_hidden_states`. Nếu là discrete, trả về các phân phối xác suất cho các latent pixels chưa bị nhiễu.

Đầu ra của [Transformer2DModel](/docs/diffusers/main/vi/api/models/transformer2d#diffusers.Transformer2DModel).