Diffusers documentation
SanaWMLTX2RefinerTransformer3DModel
SanaWMLTX2RefinerTransformer3DModel
SanaWMPipelineのステージ2として使用される、SanaWMLTX2Refinerによって駆動されるチャンク因果オートレグレッシブリファイナーDiTです。
LTX2VideoTransformer3DModel とアーキテクチャが完全に同一です。設定引数、サブモジュール、パラメータ名がすべて同じであるため、公開されている LTX-2 チェックポイントは変更なしで読み込むことができます。フォワードパスだけが異なります。
- 動画ストリームのみが実行されます（音声および音声/動画のクロスアテンションブランチはスキップされます）。
- self-attention は、attention sink と最近の精錬済み履歴を保持する明示的なスライディングウィンドウ KV キャッシュ（
SanaWMRefinerKVCache）に対して実行されるため、ブロックごとの計算量は有界であり、精錬の総コストは動画の長さに線形に比例します。 - 呼び出し元がビデオ RoPE を供給することで、各自己回帰ウィンドウがソースビデオ内の各フレームの絶対インデックスを保持できます（SanaWMLTX2RefinerTransformer3DModel.build_rotary_emb_for_absolute_positions() を参照）。
モデルは以下のように読み込むことができます:
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
class diffusers.SanaWMLTX2RefinerTransformer3DModel
< source >( in_channels: int = 128out_channels: int | None = 128patch_size: int = 1patch_size_t: int = 1num_attention_heads: int = 32attention_head_dim: int = 128cross_attention_dim: int = 4096vae_scale_factors: tuple[int, int, int] = (8, 32, 32)pos_embed_max_pos: int = 20base_height: int = 2048base_width: int = 2048gated_attn: bool = Falsecross_attn_mod: bool = Falseaudio_in_channels: int = 128audio_out_channels: int | None = 128audio_patch_size: int = 1audio_patch_size_t: int = 1audio_num_attention_heads: int = 32audio_attention_head_dim: int = 64audio_cross_attention_dim: int = 2048audio_scale_factor: int = 4audio_pos_embed_max_pos: int = 20audio_sampling_rate: int = 16000audio_hop_length: int = 160audio_gated_attn: bool = Falseaudio_cross_attn_mod: bool = Falsenum_layers: int = 48activation_fn: str = 'gelu-approximate'qk_norm: str = 'rms_norm_across_heads'norm_elementwise_affine: bool = Falsenorm_eps: float = 1e-06caption_channels: int = 3840attention_bias: bool = Trueattention_out_bias: bool = Truerope_theta: float = 10000.0rope_double_precision: bool = Truecausal_offset: int = 1timestep_scale_multiplier: int = 1000cross_attn_timestep_scale_multiplier: int = 1000rope_type: str = 'interleaved'use_prompt_embeddings = Trueperturbed_attn: bool = Falseff_bias: bool = Trueaudio_ff_bias: bool = Trueuse_prompt_adaln_single: bool = Trueuse_keyframes_abs_pos_embedding: bool = False )
パラメータ
- 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"）。
SANA-WMのステージ2として使用される、チャンク因果オートレグレッシブリファイナーTransformer。
LTX2VideoTransformer3DModel とアーキテクチャが完全に同一です。設定引数、サブモジュール、パラメータ名がすべて同じであるため、リリース済みの LTX-2 チェックポイントは変更なしで読み込むことができます。異なるのはフォワードパスです。
- 動画ストリームのみが実行されます（音声および音声/動画のクロスアテンションブランチはスキップされます）。
- self-attention は、attention sink と最近の精錬済み履歴を保持する明示的なスライディングウィンドウ KV キャッシュ（
SanaWMRefinerKVCache）に対して実行されるため、精錬のコストは AR ブロックごとに上限が設定され、動画の長さに比例して線形に増加します。 - 呼び出し元がビデオ RoPE を供給することで、各 AR ウィンドウはソースビデオ内の各フレームの絶対インデックスを保持できます（build_rotary_emb_for_absolute_positions() を参照）。
残りの引数は、LTX2VideoTransformer3DModel と一対一で対応しています。オーディオ側の引数とサブモジュールは、チェックポイントのオーディオ重みがラウンドトリップできるように保持されているだけで、リファイナのフォワードパスでは使用されません。
build_rotary_emb_for_absolute_positions
< source >( batch_size: intframe_positions: list[int]height: intwidth: intdevice: devicefps: float ) → tuple[torch.Tensor, torch.Tensor]
明示的な絶対的な潜在フレームのインデックスのリストに対して、動画用のRoPEを構築します。
SanaWMLTX2AudioVideoRotaryPosEmbed.prepare_video_coords は連続した torch.arange(num_frames) を前提としています。これは双方向推論では問題ありません。一方、スライディングウィンドウ方式の AR リファイナでは、ソース動画における各フレームの絶対インデックスを保持する必要があるため、[sink + recent + active] ウィンドウ全体で RoPE が正しい時間位相を捕捉できるようにします。
forward
< source >( hidden_states: Tensorencoder_hidden_states: Tensortimestep: Tensorvideo_rotary_emb: tuple[torch.Tensor, torch.Tensor]encoder_attention_mask: typing.Optional[torch.Tensor] = Nonekv_cache: diffusers.models.transformers.transformer_sana_wm_refiner.SanaWMRefinerKVCache | None = Nonekv_cache_mode: str | None = Noneattention_kwargs: dict | None = Nonereturn_dict: bool = True ) → Transformer2DModelOutput or tuple
パラメータ
- 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]) — アクティブなブロックの絶対フレーム位置に対する(cos, sin)RoPE です。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_peRoPE を格納します。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 にアテンションを適用し、このブロックの post-RoPE K/V をキャッシュに記録して、履歴に追加できるようにします。
Noneの場合、ブロックは現在のトークンに対して通常のフルセルフアテンションのみを実行します。 - attention_kwargs (
dict、省略可) — LoRA スケール処理に転送される任意の kwargs。 - return_dict (
bool、任意、デフォルトはTrue) —Trueの場合、通常のタプルではなく Transformer2DModelOutput を返します。
戻り値
Transformer2DModelOutput または tuple
アクティブブロックの予測速度。形状は(batch_size, num_video_tokens, out_channels)です。
単一のARブロックに対する、動画のみのフォワードパス。
SanaWMRefinerKVCache
class diffusers.models.transformers.transformer_sana_wm_refiner.SanaWMRefinerKVCache
< source >( num_layers: int )
各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
class diffusers.models.transformers.transformer_sana_wm_refiner.SanaWMRefinerKVLayerCache
< source >( )
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 の生 latents から一度だけ取得されます。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によって書き込まれるリードバック用スロット。
最後の forward で取得した post-RoPE K/V をポップします。
最後の forward で取得した pre-RoPE の K/V を pop します。
RoPE適用後の直近履歴のK/Vを返します。空の場合はNoneを返します。
RoPE 適用前の sink K/V を返します。キャプチャされていない場合（または空の場合）は None を返します。
現在のフォワードで生成されたRoPE適用後のK/Vを保存します。
現在のフォワードで生成されたRoPE適用前のK/Vを保存します。
RoPE適用後の直近履歴のK/Vを保存します。
RoPE 適用前の sink K/V を保存します。
Transformer2DModelOutput
クラス diffusers.models.modeling_outputs.Transformer2DModelOutput
< source >( sample: torch.Tensor )
パラメータ
- sample (
torch.Tensor、形状(batch_size, num_channels, height, width)、または Transformer2DModel が離散型の場合(batch size, num_vector_embeds - 1, num_latent_pixels)) —encoder_hidden_states入力を条件とする隠れ状態の出力です。離散型の場合、ノイズ除去されていない潜在ピクセルの確率分布を返します。
Transformer2DModelの出力です。