Hugging Face's logo

Diffusers documentation

SanaWMLTX2RefinerTransformer3DModel

Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

SanaWMLTX2RefinerTransformer3DModel

Refiner DiT autoregresif kausal-chunk yang digunakan sebagai tahap 2 dari SanaWMPipeline, digerakkan oleh SanaWMLTX2Refiner.

Secara arsitektur, model ini identik dengan LTX2VideoTransformer3DModel — argumen konfigurasi yang sama, submodul yang sama, nama parameter yang sama — sehingga checkpoint LTX-2 yang dirilis dapat dimuat ke dalamnya tanpa perubahan. Jalur forward-nya berbeda:

  • hanya stream video yang dijalankan (cabang cross-attention audio dan audio/video dilewati),
  • self-attention berjalan terhadap KV cache sliding-window eksplisit (SanaWMRefinerKVCache) yang menyimpan attention sink serta riwayat refined terbaru, sehingga komputasi per blok dibatasi dan total biaya refinement berskala linear dengan panjang video,
  • pemanggil menyediakan video RoPE, yang memungkinkan setiap jendela autoregresif mempertahankan indeks absolut setiap frame dalam video sumber (lihat SanaWMLTX2RefinerTransformer3DModel.build_rotary_emb_for_absolute_positions()).

Model dapat dimuat dengan:

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

kelas diffusers.SanaWMLTX2RefinerTransformer3DModel

< >

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

Parameter

  • in_channels (int, defaultnya 128) — Jumlah channel pada input.
  • out_channels (int, defaultnya 128) — Jumlah channel pada output.
  • patch_size (int, defaultnya 1) — Ukuran patch spasial yang akan digunakan di lapisan patch embedding.
  • patch_size_t (int, defaultnya 1) — Ukuran patch temporal yang akan digunakan di lapisan patch embedding.
  • num_attention_heads (int, defaultnya 32) — Jumlah head yang digunakan untuk multi-head attention.
  • attention_head_dim (int, defaultnya 128) — Jumlah saluran di setiap head.
  • cross_attention_dim (int, defaultnya 4096) — Jumlah saluran untuk kepala cross attention.
  • num_layers (int, defaultnya 48) — Jumlah lapisan blok Transformer yang akan digunakan.
  • activation_fn (str, defaultnya "gelu-approximate") — Fungsi aktivasi yang digunakan dalam feed-forward.
  • qk_norm (str, default "rms_norm_across_heads") — Lapisan normalisasi yang digunakan.
  • rope_type (str, defaultnya "interleaved") — Aplikasi RoPE mana yang akan digunakan ("interleaved" atau "split").

Transformer autoregresif penyempurna kausal-chunk yang digunakan sebagai tahap 2 SANA-WM.

Secara arsitektur identik dengan LTX2VideoTransformer3DModel — argumen config yang sama, submodule yang sama, nama parameter yang sama — sehingga checkpoint LTX-2 yang dirilis dapat dimuat ke dalamnya tanpa perubahan. Yang berbeda adalah forward pass:

  • hanya stream video yang dijalankan (cabang cross-attention audio dan audio/video dilewati),
  • self-attention berjalan terhadap cache KV sliding-window eksplisit (SanaWMRefinerKVCache) yang menyimpan attention sink serta riwayat refined terbaru, sehingga biaya refinement dibatasi per blok AR dan berskala linear dengan panjang video,
  • pemanggil menyediakan RoPE video, yang memungkinkan setiap jendela AR mempertahankan indeks absolut setiap frame dalam video sumber (lihat build_rotary_emb_for_absolute_positions()).

Argumen yang tersisa mencerminkan LTX2VideoTransformer3DModel satu per satu. Argumen dan submodul sisi audio dipertahankan murni agar bobot audio checkpoint dapat diproses bolak-balik; mereka tidak digunakan oleh forward refiner.

build_rotary_emb_for_absolute_positions

< >

( batch_size: intframe_positions: list[int]height: intwidth: intdevice: devicefps: float ) → tuple[torch.Tensor, torch.Tensor]

Parameter

  • batch_size (int) — Ukuran batch untuk menyiarkan koordinat.
  • frame_positions (list[int]) — Indeks frame laten absolut yang dicakup oleh jendela ini.
  • height (int), width (int) — Resolusi spasial laten.
  • device (torch.device) — Perangkat untuk membangun koordinat di atasnya.
  • fps (float) — Laju bingkai video, yang menggerakkan penskalaan RoPE temporal LTX-2.

Mengembalikan

tuple[torch.Tensor, torch.Tensor]

tuple RoPE (cos, sin).

Bangun video RoPE untuk daftar eksplisit indeks frame laten absolut.

SanaWMLTX2AudioVideoRotaryPosEmbed.prepare_video_coords mengasumsikan torch.arange(num_frames) yang berurutan, yang tidak masalah untuk inferensi bidireksional. Refiner AR sliding-window justru perlu mempertahankan indeks absolut setiap frame dalam video sumber, sehingga RoPE menangkap fase temporal yang benar di seluruh jendela [sink + recent + active].

forward

< >

( 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

Parameter

  • hidden_states (torch.Tensor) — Latens video yang telah dipatch dari blok aktif, dengan bentuk (batch_size, num_video_tokens, in_channels).
  • encoder_hidden_states (torch.Tensor) — Embedding teks dengan bentuk (batch_size, text_seq_len, caption_channels).
  • timestep (torch.Tensor) — Timestep dengan bentuk (batch_size, num_video_tokens), yang sudah diskalakan oleh self.config.timestep_scale_multiplier.
  • video_rotary_emb (tuple[torch.Tensor, torch.Tensor]) — (cos, sin) RoPE untuk posisi frame absolut dari blok aktif, sebagaimana dikembalikan oleh build_rotary_emb_for_absolute_positions().
  • encoder_attention_mask (torch.Tensor, opsional) — Mask perhatian teks multiplikatif dengan bentuk (batch_size, text_seq_len).
  • kv_cache (SanaWMRefinerKVCache, opsional) — KV cache sliding-window yang menyimpan attention sink per-layer dan riwayat refined terbaru, serta sink_pe RoPE bersama. Wajib diisi setiap kali kv_cache_mode disetel.
  • kv_cache_mode (str, opsional) — Salah satu dari:
    • "inject": perhatikan K/V [sink + history + current] (langkah-langkah denoising).
    • "capture_pre_rope": tanpa awalan; catat K/V pre-RoPE dari forward ini ke dalam cache (digunakan sekali untuk menabur benih attention sink dari latens stage-1 mentah).
    • "inject_and_capture_post_rope": perhatikan [sink + history + current] K/V dan catat K/V post-RoPE blok ini ke dalam cache agar dapat ditambahkan ke history.

    Ketika None, blok menjalankan full self-attention biasa hanya pada token saat ini.

  • attention_kwargs (dict, opsional) — Kwargs opsional yang diteruskan ke penanganan skala LoRA.
  • return_dict (bool, opsional, default True) — Apakah akan mengembalikan Transformer2DModelOutput alih-alih tuple biasa.

Mengembalikan

Transformer2DModelOutput atau tuple

kecepatan yang diprediksi untuk blok aktif, dengan bentuk (batch_size, num_video_tokens, out_channels).

Lewatan maju hanya video atas satu blok AR.

SanaWMRefinerKVCache

kelas diffusers.models.transformers.transformer_sana_wm_refiner.SanaWMRefinerKVCache

< >

( num_layers: int )

Parameter

  • num_layers (int) — Jumlah blok transformer untuk mengalokasikan cache per lapisan.

Wadah yang menyimpan satu [SanaWMRefinerKVLayerCache] per blok transformer, ditambah sink RoPE bersama.

Ini mengimplementasikan kontrak KV-cache rf_shifted_sink yang digunakan saat pelatihan refiner stage-2 SANA-WM. Penyempurnaan bersifat chunk-causal: block_size frame laten didenoise sekaligus sambil memperhatikan jendela terbatas dari [attention sink + riwayat terbaru + blok aktif] K/V.

  • sink_pe: tuple RoPE (cos, sin) untuk frame sink, dibangun ulang per jendela AR pada sink_rope_offset yang bergeser sehingga sink berada tepat sebelum cache kerja yang terbatas. Dibagikan antar lapisan karena RoPE tidak bergantung pada lapisan.

SanaWMRefinerKVLayerCache

kelas diffusers.models.transformers.transformer_sana_wm_refiner.SanaWMRefinerKVLayerCache

< >

( )

Cache KV per layer untuk refiner AR chunk-causal tahap-2 SANA-WM.

Menyimpan dua bagian dari prefix sliding-window yang menjadi fokus self-attention refiner, ditambah slot untuk membaca kembali K/V yang ditangkap oleh forward terakhir. Semua tensor berukuran (batch_size, num_tokens, inner_dim) (yaitu sebelum unflatten per-head), sesuai dengan layout yang digabungkan oleh self-attention refiner.

  • sink_k_pre / sink_v: pre-RoPE K/V dari frame attention-sink, diambil sekali dari latens stage-1 mentah. Nilai ini disimpan dalam bentuk pre-RoPE sehingga setiap jendela AR dapat menerapkan ulang RoPE pada offset sink yang digeser sesuai posisinya (SanaWMRefinerKVCache.sink_pe).
  • history_k / history_v: post-RoPE K/V dari frame-frame terbaru yang sudah disempurnakan, siap untuk digabungkan apa adanya.
  • captured_k_pre / captured_v_pre dan captured_k_post / captured_v_post: slot readback yang ditulis oleh kv_cache_mode capture.

get_captured_post_rope

< >

( )

Ambil K/V pasca-RoPE yang ditangkap oleh forward terakhir.

get_captured_pre_rope

< >

( )

Ambil K/V pra-RoPE yang ditangkap oleh forward terakhir.

get_history

< >

( )

Kembalikan K/V riwayat-recent pasca-RoPE, atau None jika kosong.

get_sink

< >

( )

Kembalikan sink K/V pre-RoPE, atau None jika belum ditangkap (atau kosong).

store_captured_post_rope

< >

( key: Tensorvalue: Tensor )

Simpan K/V pasca-RoPE yang dihasilkan oleh forward saat ini.

store_captured_pre_rope

< >

( key: Tensorvalue: Tensor )

Simpan K/V pra-RoPE yang dihasilkan oleh forward saat ini.

store_history

< >

( history_k: Tensorhistory_v: Tensor )

Simpan K/V riwayat terbaru pasca-RoPE.

store_sink

< >

( sink_k_pre: Tensorsink_v: Tensor )

Simpan K/V sink pra-RoPE.

Transformer2DModelOutput

kelas diffusers.models.modeling_outputs.Transformer2DModelOutput

< >

( sample: torch.Tensor )

Parameter

  • sample (torch.Tensor dengan bentuk (batch_size, num_channels, height, width) atau (batch size, num_vector_embeds - 1, num_latent_pixels) jika Transformer2DModel bersifat diskrit) — Output hidden states yang dikondisikan pada input encoder_hidden_states. Jika diskrit, mengembalikan distribusi probabilitas untuk latent pixels yang tidak berisik.

Output dari Transformer2DModel.

Update on GitHub

←SanaVideoTransformer3DModel SanaWMTransformer3DModel→