# SanaWMLTX2RefinerTransformer3DModel

O refiner DiT autoregressivo causal por blocos usado como etapa 2 do [SanaWMPipeline](/docs/diffusers/v0.43.0/pt/api/pipelines/sana_wm#diffusers.SanaWMPipeline), dirigido por [SanaWMLTX2Refiner](/docs/diffusers/v0.43.0/pt/api/pipelines/sana_wm#diffusers.SanaWMLTX2Refiner).

Ele é arquitetonicamente idêntico ao [LTX2VideoTransformer3DModel](/docs/diffusers/v0.43.0/pt/api/models/ltx2_video_transformer3d#diffusers.LTX2VideoTransformer3DModel) — mesmos argumentos de configuração, mesmos submódulos, mesmos nomes de parâmetros — de modo que um checkpoint LTX-2 publicado é carregado nele sem alterações. A passagem para frente difere:

* somente o fluxo de vídeo é executado (os ramos de atenção cruzada de áudio e áudio/vídeo são ignorados),
* o self-attention é executado contra um cache KV de janela deslizante explícito (`SanaWMRefinerKVCache`) que armazena o sink de atenção e o histórico refinado recente, de modo que o cálculo por bloco é limitado e o custo total de refinamento escala linearmente com o comprimento do vídeo,
* o chamador fornece o RoPE do vídeo, o que permite que cada janela autoregressiva mantenha o índice absoluto de cada quadro no vídeo de origem (veja [SanaWMLTX2RefinerTransformer3DModel.build_rotary_emb_for_absolute_positions()](/docs/diffusers/v0.43.0/pt/api/models/sana_wm_refiner_transformer3d#diffusers.SanaWMLTX2RefinerTransformer3DModel.build_rotary_emb_for_absolute_positions)).

O modelo pode ser carregado com:

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

[Fonte](https://github.com/huggingface/diffusers/blob/v0.43.0/src/diffusers/models/transformers/transformer_sana_wm_refiner.py#L1110)

**Parâmetros:**

in_channels (`int`, padrão `128`) : O número de canais na entrada.

out_channels (`int`, padrão `128`) : O número de canais na saída.

patch_size (`int`, padrão `1`) : O tamanho dos patches espaciais a serem usados na camada de incorporação de patches.

patch_size_t (`int`, padrão `1`) : O tamanho dos patches temporais a serem usados na camada de incorporação de patches.

num_attention_heads (`int`, padrão `32`) : O número de cabeças a serem usadas para a atenção multi-cabeça.

attention_head_dim (`int`, padrão `128`) : O número de canais em cada head.

cross_attention_dim (`int`, padrão `4096`) : O número de canais para as cabeças de atenção cruzada.

num_layers (`int`, padrão `48`) : O número de blocos de Transformer a serem usados.

activation_fn (`str`, padrão `"gelu-approximate"`) : Função de ativação a ser usada no feed-forward.

qk_norm (`str`, padrão `"rms_norm_across_heads"`) : A camada de normalização a ser usada.

rope_type (`str`, padrão `"interleaved"`) : Qual aplicação de RoPE usar (`"interleaved"` ou `"split"`).

O transformer autoregressivo refinador causal por blocos usado como etapa 2 do SANA-WM.

Arquiteturalmente idêntico a [LTX2VideoTransformer3DModel](/docs/diffusers/v0.43.0/pt/api/models/ltx2_video_transformer3d#diffusers.LTX2VideoTransformer3DModel) — mesmos argumentos de configuração, mesmos submódulos, mesmos nomes de parâmetros — de modo que um checkpoint LTX-2 publicado é carregado nele sem alterações. O que difere é o forward pass:

* somente o fluxo de vídeo é executado (os ramos de atenção cruzada de áudio e áudio/vídeo são ignorados),
* o self-attention é executado contra um cache KV de janela deslizante explícito (`SanaWMRefinerKVCache`) que mantém o sink de atenção e o histórico refinado recente, de modo que o custo de refinamento é limitado por bloco AR e escala linearmente com o comprimento do vídeo,
* o chamador fornece o RoPE do vídeo, o que permite que cada janela AR mantenha o índice absoluto de cada quadro no vídeo de origem (veja [build_rotary_emb_for_absolute_positions()](/docs/diffusers/v0.43.0/pt/api/models/sana_wm_refiner_transformer3d#diffusers.SanaWMLTX2RefinerTransformer3DModel.build_rotary_emb_for_absolute_positions)).

Os argumentos restantes espelham os de [LTX2VideoTransformer3DModel](/docs/diffusers/v0.43.0/pt/api/models/ltx2_video_transformer3d#diffusers.LTX2VideoTransformer3DModel) um a um. Os argumentos e submódulos do lado de áudio são mantidos puramente para que os pesos de áudio do checkpoint façam ida e volta; eles não são usados pelo forward do 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)
```

[Fonte](https://github.com/huggingface/diffusers/blob/v0.43.0/src/diffusers/models/transformers/transformer_sana_wm_refiner.py#L1376)

**Parâmetros:**

batch_size (`int`) : Tamanho do lote para o qual as coordenadas devem ser transmitidas.

frame_positions (`list[int]`) : Índices absolutos de quadros latentes cobertos por esta janela.

height (`int`), largura (`int`) : Resolução espacial latente.

device (`torch.device`) : Dispositivo em que as coordenadas serão construídas.

fps (`float`) : Taxa de quadros do vídeo, que controla a escala do RoPE temporal do LTX-2.

**Retorna:** `tuple[torch.Tensor, torch.Tensor]`

a tupla RoPE `(cos, sin)`.

Construa o RoPE de vídeo para uma lista explícita de índices absolutos de quadros latentes.

`SanaWMLTX2AudioVideoRotaryPosEmbed.prepare_video_coords` pressupõe um `torch.arange(num_frames)` contíguo, o que é adequado para inferência bidirecional. O refinador AR de janela deslizante, por sua vez, precisa manter o índice absoluto de cada quadro no vídeo de origem, para que o RoPE capture a fase temporal correta ao longo da janela `[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)
```

[Fonte](https://github.com/huggingface/diffusers/blob/v0.43.0/src/diffusers/models/transformers/transformer_sana_wm_refiner.py#L1433)

**Parâmetros:**

hidden_states (`torch.Tensor`) : Latentes de vídeo em patches do bloco ativo, com forma `(batch_size, num_video_tokens, in_channels)`.

encoder_hidden_states (`torch.Tensor`) : Incorporações de texto com forma `(batch_size, text_seq_len, caption_channels)`.

timestep (`torch.Tensor`) : Timestep de forma `(batch_size, num_video_tokens)`, já escalado por `self.config.timestep_scale_multiplier`.

video_rotary_emb (`tuple[torch.Tensor, torch.Tensor]`) : O RoPE `(cos, sin)` para as posições absolutas do quadro ativo do bloco, conforme retornado por [build_rotary_emb_for_absolute_positions()](/docs/diffusers/v0.43.0/pt/api/models/sana_wm_refiner_transformer3d#diffusers.SanaWMLTX2RefinerTransformer3DModel.build_rotary_emb_for_absolute_positions).

encoder_attention_mask (`torch.Tensor`, *opcional*) : Máscara de atenção multiplicativa do texto com forma `(batch_size, text_seq_len)`.

kv_cache (`SanaWMRefinerKVCache`, *opcional*) : Cache de KV de janela deslizante que armazena o sink de atenção por camada e o histórico refinado recente, além do RoPE compartilhado `sink_pe`. Obrigatório sempre que `kv_cache_mode` estiver definido.

kv_cache_mode (`str`, *opcional*) : Um dos seguintes: - `"inject"`: atende a `[sink + history + current]` K/V (as etapas de desruido). - `"capture_pre_rope"`: sem prefixo; registra o K/V pré-RoPE deste forward no cache (usado uma vez para semear o sink de atenção a partir dos latentes brutos da etapa 1). - `"inject_and_capture_post_rope"`: atende a `[sink + history + current]` K/V e registra o K/V pós-RoPE deste bloco no cache para que possa ser anexado ao histórico. Quando `None`, o bloco executa autoatenção completa simples apenas sobre os tokens atuais.

attention_kwargs (`dict`, *opcional*) : Kwargs opcionais encaminhados para o tratamento de escala do LoRA.

return_dict (`bool`, *opcional*, padrão `True`) : Se deve retornar um [Transformer2DModelOutput](/docs/diffusers/v0.43.0/pt/api/models/cogview3plus_transformer2d#diffusers.models.modeling_outputs.Transformer2DModelOutput) em vez de uma tupla simples.

**Retorna:** [Transformer2DModelOutput](/docs/diffusers/v0.43.0/pt/api/models/cogview3plus_transformer2d#diffusers.models.modeling_outputs.Transformer2DModelOutput) ou `tuple`

a velocidade prevista para o bloco ativo, de forma `(batch_size, num_video_tokens, out_channels)`.

Passada de propagação direta apenas de vídeo sobre um único bloco 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)
```

[Fonte](https://github.com/huggingface/diffusers/blob/v0.43.0/src/diffusers/models/transformers/transformer_sana_wm_refiner.py#L745)

**Parâmetros:**

num_layers (*int*) : Número de blocos de transformer para os quais alocar um cache por camada.

Contêiner que mantém uma [*SanaWMRefinerKVLayerCache*] por bloco do transformer, além do RoPE de sink compartilhado.

Isso implementa o contrato de KV-cache `rf_shifted_sink` com o qual o refinador do estágio 2 do SANA-WM foi treinado. O refinamento é causal por blocos: *block_size* quadros latentes são desruidos de cada vez, enquanto prestam atenção a uma janela limitada de *[attention sink + histórico recente + bloco ativo]* K/V.

- `sink_pe`: a tupla RoPE *(cos, sin)* para os quadros de sink, reconstruída a cada janela de AR no `sink_rope_offset` deslizante, de modo que o sink fique imediatamente antes do cache de trabalho limitado. Compartilhada entre as camadas, pois o RoPE não depende da camada.

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

[Fonte](https://github.com/huggingface/diffusers/blob/v0.43.0/src/diffusers/models/transformers/transformer_sana_wm_refiner.py#L656)

Cache de KV por camada para o refinador AR causal por blocos do estágio 2 do SANA-WM.

Contém as duas metades do prefixo da janela deslizante ao qual a autoatenção do refiner presta atenção, além de um slot para ler de volta o K/V capturado pelo último forward. Todos os tensors são *(batch_size, num_tokens, inner_dim)* (ou seja, antes do unflatten por cabeça), correspondendo ao layout que a autoatenção do refiner concatena.

- `sink_k_pre` / `sink_v`: **pre**-RoPE K/V dos quadros de attention-sink, capturados uma única vez dos latentes brutos da etapa 1. Eles são armazenados pré-RoPE para que cada janela AR possa reaplicar o RoPE em seu próprio deslocamento de sink (`SanaWMRefinerKVCache.sink_pe`).
- `history_k` / `history_v`: K/V **post**-RoPE dos quadros recentes já refinados, prontos para serem concatenados como estão.
- `captured_k_pre` / `captured_v_pre` e `captured_k_post` / `captured_v_post`: slots de leitura gravados pelos modos de captura `kv_cache_mode`.

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

```python
get_captured_post_rope()
```

[Fonte](https://github.com/huggingface/diffusers/blob/v0.43.0/src/diffusers/models/transformers/transformer_sana_wm_refiner.py#L724)

Extrair o K/V pós-RoPE capturado pelo último forward.

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

```python
get_captured_pre_rope()
```

[Fonte](https://github.com/huggingface/diffusers/blob/v0.43.0/src/diffusers/models/transformers/transformer_sana_wm_refiner.py#L710)

Pegue o K/V pré-RoPE capturado pelo último forward.

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

```python
get_history()
```

[Fonte](https://github.com/huggingface/diffusers/blob/v0.43.0/src/diffusers/models/transformers/transformer_sana_wm_refiner.py#L699)

Retorna o K/V do histórico recente pós-RoPE, ou `None` se estiver vazio.

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

```python
get_sink()
```

[Fonte](https://github.com/huggingface/diffusers/blob/v0.43.0/src/diffusers/models/transformers/transformer_sana_wm_refiner.py#L688)

Retorna o K/V de sink pré-RoPE, ou `None` se ele não tiver sido capturado (ou estiver vazio).

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

[Fonte](https://github.com/huggingface/diffusers/blob/v0.43.0/src/diffusers/models/transformers/transformer_sana_wm_refiner.py#L719)

Armazene o K/V pós-RoPE produzido pelo forward atual.

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

[Fonte](https://github.com/huggingface/diffusers/blob/v0.43.0/src/diffusers/models/transformers/transformer_sana_wm_refiner.py#L705)

Armazene o K/V pré-RoPE produzido pelo forward atual.

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

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

[Fonte](https://github.com/huggingface/diffusers/blob/v0.43.0/src/diffusers/models/transformers/transformer_sana_wm_refiner.py#L694)

Armazene o K/V do histórico recente pós-RoPE.

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

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

[Fonte](https://github.com/huggingface/diffusers/blob/v0.43.0/src/diffusers/models/transformers/transformer_sana_wm_refiner.py#L683)

Armazene o K/V do sink pré-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)
```

[Fonte](https://github.com/huggingface/diffusers/blob/v0.43.0/src/diffusers/models/modeling_outputs.py#L21)

**Parâmetros:**

sample (`torch.Tensor` de forma `(batch_size, num_channels, height, width)` ou `(batch size, num_vector_embeds - 1, num_latent_pixels)` se [Transformer2DModel](/docs/diffusers/v0.43.0/pt/api/models/transformer2d#diffusers.Transformer2DModel) for discreto) : A saída dos estados ocultos condicionada à entrada `encoder_hidden_states`. Se discreto, retorna distribuições de probabilidade para os pixels latentes sem ruído.

A saída de [Transformer2DModel](/docs/diffusers/v0.43.0/pt/api/models/transformer2d#diffusers.Transformer2DModel).