# SanaWMLTX2RefinerTransformer3DModel

Чанково-каузальный авто-регрессивный рефинер DiT, используемый как второй этап [SanaWMPipeline](/docs/diffusers/v0.43.0/ru/api/pipelines/sana_wm#diffusers.SanaWMPipeline), управляемый [SanaWMLTX2Refiner](/docs/diffusers/v0.43.0/ru/api/pipelines/sana_wm#diffusers.SanaWMLTX2Refiner).

Архитектурно он идентичен [LTX2VideoTransformer3DModel](/docs/diffusers/v0.43.0/ru/api/models/ltx2_video_transformer3d#diffusers.LTX2VideoTransformer3DModel) — те же аргументы конфигурации, те же подмодули, те же имена параметров — поэтому опубликованный чекпоинт LTX-2 загружается в него без изменений. Прямой проход (forward pass) отличается:

* обрабатывается только видеопоток (ветви кросс-внимания аудио и аудио/видео пропускаются),
* self-attention выполняется с использованием явного скользящего окна KV-кэша (`SanaWMRefinerKVCache`), содержащего attention sink и недавнюю уточнённую историю, поэтому вычисления на блок ограничены, а общая стоимость уточнения линейно масштабируется с длиной видео,
* вызывающая сторона предоставляет видео RoPE, что позволяет каждому авто-регрессивному окну сохранять абсолютный индекс каждого кадра в исходном видео (см. [SanaWMLTX2RefinerTransformer3DModel.build_rotary_emb_for_absolute_positions()](/docs/diffusers/v0.43.0/ru/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[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)
```

[Источник](https://github.com/huggingface/diffusers/blob/v0.43.0/src/diffusers/models/transformers/transformer_sana_wm_refiner.py#L1110)

**Параметры:**

in_channels (`int`, по умолчанию `128`) : количество каналов на входе.

out_channels (`int`, по умолчанию `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`) : Количество каналов для головок кросс-внимания.

num_layers (`int`, по умолчанию `48`) : Количество блоков Transformer, используемых в модели.

activation_fn (`str`, по умолчанию `"gelu-approximate"`) : Функция активации, используемая в полносвязном слое.

qk_norm (`str`, по умолчанию `"rms_norm_across_heads"`) : Слой нормализации, который следует использовать.

rope_type (`str`, по умолчанию `"interleaved"`) : Какой способ применения RoPE использовать (`"interleaved"` или `"split"`).

Chunk-causal autoregressive refiner transformer, используемый на этапе 2 SANA-WM.

Архитектурно идентичен [LTX2VideoTransformer3DModel](/docs/diffusers/v0.43.0/ru/api/models/ltx2_video_transformer3d#diffusers.LTX2VideoTransformer3DModel) — те же аргументы конфигурации, те же подмодули, те же имена параметров — поэтому опубликованный чекпоинт LTX-2 загружается в него без изменений. Отличается только прямой проход:

* обрабатывается только видеопоток (ветви кросс-внимания аудио и аудио/видео пропускаются),
* self-attention выполняется с использованием явного скользящего окна KV-кэша (`SanaWMRefinerKVCache`), содержащего attention sink и недавнюю уточнённую историю, поэтому стоимость уточнения ограничена для каждого AR-блока и линейно масштабируется с длиной видео,
* вызывающая сторона предоставляет видео RoPE, что позволяет каждому окну AR сохранять абсолютный индекс каждого кадра в исходном видео (см. [build_rotary_emb_for_absolute_positions()](/docs/diffusers/v0.43.0/ru/api/models/sana_wm_refiner_transformer3d#diffusers.SanaWMLTX2RefinerTransformer3DModel.build_rotary_emb_for_absolute_positions)).

Остальные аргументы полностью соответствуют аргументам [LTX2VideoTransformer3DModel](/docs/diffusers/v0.43.0/ru/api/models/ltx2_video_transformer3d#diffusers.LTX2VideoTransformer3DModel). Аргументы и подмодули, связанные с аудио, сохранены исключительно для того, чтобы аудио-веса контрольной точки корректно восстанавливались; они не используются в прямом проходе рефинера.

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

[Источник](https://github.com/huggingface/diffusers/blob/v0.43.0/src/diffusers/models/transformers/transformer_sana_wm_refiner.py#L1376)

**Параметры:**

batch_size (`int`) : Размер батча, к которому нужно распространить координаты.

frame_positions (`list[int]`) : Абсолютные индексы латентных кадров, охватываемых этим окном.

height (`int`), ширина (`int`) : Пространственное разрешение латентовых представлений.

device (`torch.device`) : Устройство, на котором будут построены координаты.

fps (`float`) : Частота кадров видео, определяющая масштабирование временного RoPE в LTX-2.

**Возвращает:** `tuple[torch.Tensor, torch.Tensor]`

кортеж RoPE `(cos, sin)`.

Создайте видео-RoPE для явного списка абсолютных индексов латентных кадров.

`SanaWMLTX2AudioVideoRotaryPosEmbed.prepare_video_coords` предполагает непрерывный `torch.arange(num_frames)`, что подходит для двунаправленного вывода. Уточнитель с AR-окном скользящего типа, напротив, должен сохранять абсолютный индекс каждого кадра в исходном видео, чтобы RoPE корректно учитывал временную фазу в окне `[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)
```

[Источник](https://github.com/huggingface/diffusers/blob/v0.43.0/src/diffusers/models/transformers/transformer_sana_wm_refiner.py#L1433)

**Параметры:**

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]`) : RoPE `(cos, sin)` для абсолютных позиций кадров активного блока, как возвращено функцией [build_rotary_emb_for_absolute_positions()](/docs/diffusers/v0.43.0/ru/api/models/sana_wm_refiner_transformer3d#diffusers.SanaWMLTX2RefinerTransformer3DModel.build_rotary_emb_for_absolute_positions).

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"`: внимание к K/V `[sink + history + current]` (шаги денойзинга). - `"capture_pre_rope"`: без префикса; записать K/V до применения RoPE для этого прохода в кэш (используется один раз для инициализации attention sink из исходных латентов первого этапа). - `"inject_and_capture_post_rope"`: внимание к K/V `[sink + history + current]` и запись K/V после применения RoPE для этого блока в кэш, чтобы их можно было добавить в историю.  Если `None`, блок выполняет обычное полное само-внимание только по текущим токенам.

attention_kwargs (`dict`, *необязательный*) : Необязательные kwargs, передаваемые в обработку масштаба LoRA.

return_dict (`bool`, *необязательный*, по умолчанию `True`) : следует ли возвращать [Transformer2DModelOutput](/docs/diffusers/v0.43.0/ru/api/models/cogview3plus_transformer2d#diffusers.models.modeling_outputs.Transformer2DModelOutput) вместо обычного кортежа.

**Возвращает:** [Transformer2DModelOutput](/docs/diffusers/v0.43.0/ru/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/v0.43.0/src/diffusers/models/transformers/transformer_sana_wm_refiner.py#L745)

**Параметры:**

num_layers (*int*) : Количество блоков трансформера, для которых выделяется кэш на каждый слой.

Контейнер, содержащий по одному *SanaWMRefinerKVLayerCache* на каждый блок трансформера, а также общий sink RoPE.

Это реализует контракт KV-кэша `rf_shifted_sink`, с которым обучался рефинер второго этапа SANA-WM. Рефинирование является чанково-каузальным: *block_size* латентных кадров денойзируются за один раз, при этом внимание обращается к ограниченному окну *[attention sink + recent history + active block]* K/V.

- `sink_pe`: кортеж *(cos, sin)* RoPE для sink-кадров, пересобираемый для каждого 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/v0.43.0/src/diffusers/models/transformers/transformer_sana_wm_refiner.py#L656)

Кэш KV для каждого слоя для этапа 2 SANA-WM, использующего чанк-каузальный AR-рефинер.

Хранит две половины префикса скользящего окна, на которые обращается само-внимание рефинера, а также слот для чтения обратно K/V, захваченных последним forward. Все тензоры имеют размерность *(batch_size, num_tokens, inner_dim)* (то есть до unflatten по головам), что соответствует раскладке, которую само-внимание рефинера использует при конкатенации.

- `sink_k_pre` / `sink_v`: **pre**-RoPE K/V для кадров attention-sink, захваченные один раз из исходных латентов первого этапа. Они хранятся в формате pre-RoPE, чтобы каждое AR-окно могло повторно применить RoPE со своим сдвинутым смещением sink (`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/v0.43.0/src/diffusers/models/transformers/transformer_sana_wm_refiner.py#L724)

Снимите K/V после RoPE, полученные последним прямым проходом.

#### 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/v0.43.0/src/diffusers/models/transformers/transformer_sana_wm_refiner.py#L710)

Достаньте K/V до применения RoPE, полученные в последнем прямом проходе.

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

```python
get_history()
```

[Источник](https://github.com/huggingface/diffusers/blob/v0.43.0/src/diffusers/models/transformers/transformer_sana_wm_refiner.py#L699)

Возвращает K/V недавней истории после применения RoPE или `None`, если пусто.

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

```python
get_sink()
```

[Источник](https://github.com/huggingface/diffusers/blob/v0.43.0/src/diffusers/models/transformers/transformer_sana_wm_refiner.py#L688)

Возвращает sink K/V до применения RoPE или `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/v0.43.0/src/diffusers/models/transformers/transformer_sana_wm_refiner.py#L719)

Сохраните K/V после применения RoPE, полученные в текущем прямом проходе.

#### 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/v0.43.0/src/diffusers/models/transformers/transformer_sana_wm_refiner.py#L705)

Сохраните K/V до применения RoPE, полученные в текущем прямом проходе.

#### 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/v0.43.0/src/diffusers/models/transformers/transformer_sana_wm_refiner.py#L694)

Сохраните K/V недавней истории после применения RoPE.

#### 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/v0.43.0/src/diffusers/models/transformers/transformer_sana_wm_refiner.py#L683)

Сохраните K/V sink до применения 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)
```

[Источник](https://github.com/huggingface/diffusers/blob/v0.43.0/src/diffusers/models/modeling_outputs.py#L21)

**Параметры:**

sample (`torch.Tensor` формы `(batch_size, num_channels, height, width)` или `(batch size, num_vector_embeds - 1, num_latent_pixels)`, если [Transformer2DModel](/docs/diffusers/v0.43.0/ru/api/models/transformer2d#diffusers.Transformer2DModel) дискретный): выходные скрытые состояния, обусловленные входом `encoder_hidden_states`. Если модель дискретная, возвращает распределения вероятностей для ненашумленных латентных пикселей.

Вывод [Transformer2DModel](/docs/diffusers/v0.43.0/ru/api/models/transformer2d#diffusers.Transformer2DModel).