Hugging Face's logo

Diffusers documentation

SanaWMLTX2RefinerTransformer3DModel

You are viewing main version, which requires installation from source. If you'd like regular pip install, checkout the latest stable version (v0.43.0).
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→