# SanaWMLTX2RefinerTransformer3DModel

[SanaWMPipeline](/docs/diffusers/main/ja/api/pipelines/sana_wm#diffusers.SanaWMPipeline)のステージ2として使用される、[SanaWMLTX2Refiner](/docs/diffusers/main/ja/api/pipelines/sana_wm#diffusers.SanaWMLTX2Refiner)によって駆動されるチャンク因果オートレグレッシブリファイナーDiTです。

[LTX2VideoTransformer3DModel](/docs/diffusers/main/ja/api/models/ltx2_video_transformer3d#diffusers.LTX2VideoTransformer3DModel) とアーキテクチャが完全に同一です。設定引数、サブモジュール、パラメータ名がすべて同じであるため、公開されている LTX-2 チェックポイントは変更なしで読み込むことができます。フォワードパスだけが異なります。

* 動画ストリームのみが実行されます（音声および音声/動画のクロスアテンションブランチはスキップされます）。
* self-attention は、attention sink と最近の精錬済み履歴を保持する明示的なスライディングウィンドウ KV キャッシュ（`SanaWMRefinerKVCache`）に対して実行されるため、ブロックごとの計算量は有界であり、精錬の総コストは動画の長さに比例して線形に増加します。
* 呼び出し元がビデオ RoPE を供給し、これにより各自己回帰ウィンドウはソースビデオ内の各フレームの絶対インデックスを保持できます（[SanaWMLTX2RefinerTransformer3DModel.build_rotary_emb_for_absolute_positions()](/docs/diffusers/main/ja/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/main/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"`) : feed-forward で使用する活性化関数。

qk_norm (`str`、デフォルトは`"rms_norm_across_heads"`) : 使用する正規化レイヤー。

rope_type (`str`、デフォルトは `"interleaved"`) : 使用する RoPE の適用方法（`"interleaved"` または `"split"`）。

SANA-WMのステージ2として使用される、チャンク因果オートレグレッシブリファイナーTransformer。

[LTX2VideoTransformer3DModel](/docs/diffusers/main/ja/api/models/ltx2_video_transformer3d#diffusers.LTX2VideoTransformer3DModel) とアーキテクチャが完全に同一です。設定引数、サブモジュール、パラメータ名がすべて同じであるため、公開されている LTX-2 チェックポイントは変更なしで読み込むことができます。異なるのはフォワードパスです。

* 動画ストリームのみが実行されます（音声および音声/動画のクロスアテンションブランチはスキップされます）。
* self-attention は、attention sink と最近の精錬済み履歴を保持する明示的なスライディングウィンドウ KV キャッシュ（`SanaWMRefinerKVCache`）に対して実行されるため、精錬のコストは AR ブロックごとに上限が設定され、動画の長さに比例して線形に増加します。
* 呼び出し元が動画用の RoPE を供給し、これにより各 AR ウィンドウはソース動画における各フレームの絶対インデックスを保持できます（[build_rotary_emb_for_absolute_positions()](/docs/diffusers/main/ja/api/models/sana_wm_refiner_transformer3d#diffusers.SanaWMLTX2RefinerTransformer3DModel.build_rotary_emb_for_absolute_positions) を参照）。

残りの引数は、[LTX2VideoTransformer3DModel](/docs/diffusers/main/ja/api/models/ltx2_video_transformer3d#diffusers.LTX2VideoTransformer3DModel) と一対一で対応しています。オーディオ側の引数とサブモジュールは、チェックポイントのオーディオ重みがラウンドトリップできるように保持されているだけで、リファイナのフォワードパスでは使用されません。

#### 絶対位置用のロータリー埋め込みを構築[[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/main/src/diffusers/models/transformers/transformer_sana_wm_refiner.py#L1376)

**パラメータ:**

batch_size (`int`) : 座標をブロードキャストするバッチサイズ。

frame_positions (`list[int]`) : このウィンドウがカバーする絶対的な潜在フレームのインデックス。

height (`int`)、width (`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)` を前提としています。これは双方向推論では問題ありません。一方、スライディングウィンドウ方式の AR リファイナは、ソース動画における各フレームの絶対インデックスを保持する必要があるため、`[sink + recent + active]` ウィンドウ全体で RoPE が正しい時間位相を捕捉できるようにします。

#### 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/main/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`) : `self.config.timestep_scale_multiplier` でスケーリング済みの、`(batch_size, num_video_tokens)` の形状を持つタイムステップ。

video_rotary_emb (`tuple[torch.Tensor, torch.Tensor]`) : アクティブなブロックの絶対フレーム位置に対する `(cos, sin)` RoPE。 [build_rotary_emb_for_absolute_positions()](/docs/diffusers/main/ja/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`、*任意*) : 各レイヤーの attention sink と最近の refined history を保持するスライディングウィンドウ KV キャッシュ、および共有の `sink_pe` RoPE。`kv_cache_mode` が設定されている場合は必須です。

kv_cache_mode (`str`, *任意*) : 以下のいずれかです: - `"inject"`: `[sink + history + current]` の K/V にアテンションを適用します（デノイジングステップ）。 - `"capture_pre_rope"`: プレフィックスなし。このフォワードの RoPE 適用前の K/V をキャッシュに記録します（ステージ1の生ラテンツからアテンションシンクをシードするために一度だけ使用）。 - `"inject_and_capture_post_rope"`: `[sink + history + current]` の K/V にアテンションを適用し、このブロックの RoPE 適用後の K/V をキャッシュに記録して、履歴に追加できるようにします。 `None` の場合、このブロックは現在のトークンに対して通常のフルセルフアテンションのみを実行します。

attention_kwargs (`dict`, *任意*) : LoRA スケール処理に転送される任意の kwargs。

return_dict (`bool`、*任意*、デフォルトは`True`) : プレーンなタプルではなく、[Transformer2DModelOutput](/docs/diffusers/main/ja/api/models/cogview3plus_transformer2d#diffusers.models.modeling_outputs.Transformer2DModelOutput)を返すかどうか。

**戻り値:** [Transformer2DModelOutput](/docs/diffusers/main/ja/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#L745)

**パラメータ:**

num_layers (*int*) : 層ごとにキャッシュを割り当てるための Transformer ブロックの数。

各transformerブロックごとに1つの*SanaWMRefinerKVLayerCache*を保持するコンテナ、および共有のsink RoPE。

これは、SANA-WM のステージ2リファイナが学習時に使用していた `rf_shifted_sink` KVキャッシュの契約を実装しています。リファインメントはチャンク因果的であり、*block_size* 個の潜在フレームを一度にデノイズしながら、*[アテンションシンク + 直近の履歴 + アクティブブロック]* という有界ウィンドウの K/V にアテンションを行います。

- `sink_pe`: sink フレーム用の *(cos, sin)* RoPE タプル。スライディングする `sink_rope_offset` に対して AR ウィンドウごとに再構築され、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#L656)

SANA-WM の stage-2 chunk-causal AR refiner における、レイヤーごとの KV キャッシュ。

リファイナーの自己注意機構が参照するスライディングウィンドウのプレフィックスの2つの半分のテンソルを保持し、さらに最後のフォワードでキャプチャされた K/V を読み戻すためのスロットも備えています。すべてのテンソルは *(batch_size, num_tokens, inner_dim)*（つまり、ヘッドごとのアンフラットニング前のレイアウト）であり、リファイナーの自己注意機構が連結するレイアウトと一致しています。

- `sink_k_pre` / `sink_v`: **pre**-RoPE 適用前の attention-sink フレームの K/V です。stage-1 の生ラテンツから一度だけ取得されます。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#L724)

最後の forward で取得した post-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#L710)

最後の forward で取得した pre-RoPE の K/V を pop します。

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

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

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

現在のフォワードで生成された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#L705)

現在のフォワードで生成された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#L694)

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

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`、shape は `(batch_size, num_channels, height, width)`、または [Transformer2DModel](/docs/diffusers/main/ja/api/models/transformer2d#diffusers.Transformer2DModel) が離散型の場合 `(batch size, num_vector_embeds - 1, num_latent_pixels)`) : `encoder_hidden_states` 入力を条件とする隠れ状態の出力です。離散型の場合、ノイズ除去されていない潜在ピクセルの確率分布を返します。

[Transformer2DModel](/docs/diffusers/main/ja/api/models/transformer2d#diffusers.Transformer2DModel)の出力です。