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

DiT tự hồi quy nhân quả theo chunk (chunk-causal autoregressive refiner) được sử dụng làm giai đoạn 2 của SanaWMPipeline, được điều khiển bởi SanaWMLTX2Refiner.

Về mặt kiến trúc, nó giống hệt LTX2VideoTransformer3DModel — cùng các đối số cấu hình, cùng các module con, cùng tên tham số — do đó một checkpoint LTX-2 đã phát hành có thể được nạp vào mà không cần thay đổi. Quá trình truyền tiến (forward pass) khác nhau:

  • chỉ luồng video được chạy (các nhánh chú ý chéo giữa âm thanh và âm thanh/video bị bỏ qua),
  • self-attention chạy trên một KV cache trượt cửa sổ tường minh (SanaWMRefinerKVCache) chứa attention sink cùng lịch sử tinh chỉnh gần nhất, do đó chi phí tính toán cho mỗi khối bị giới hạn và tổng chi phí tinh chỉnh tăng tuyến tính theo độ dài video,
  • người gọi cung cấp video RoPE, cho phép mỗi cửa sổ tự hồi quy giữ lại chỉ số tuyệt đối của từng khung hình trong video nguồn (xem SanaWMLTX2RefinerTransformer3DModel.build_rotary_emb_for_absolute_positions()).

Mô hình có thể được tải bằng cách:

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

lớp 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 )

Tham số

  • in_channels (int, mặc định là 128) — Số lượng kênh trong đầu vào.
  • out_channels (int, mặc định là 128) — Số lượng kênh trong đầu ra.
  • patch_size (int, mặc định là 1) — Kích thước của các patch không gian được sử dụng trong lớp nhúng patch.
  • patch_size_t (int, mặc định là 1) — Kích thước của các patch thời gian được sử dụng trong lớp nhúng patch.
  • num_attention_heads (int, mặc định là 32) — Số lượng đầu (heads) được sử dụng cho cơ chế chú ý đa đầu (multi-head attention).
  • attention_head_dim (int, mặc định là 128) — Số lượng kênh trong mỗi head.
  • cross_attention_dim (int, mặc định là 4096) — Số lượng kênh cho các đầu chú ý chéo (cross attention).
  • num_layers (int, mặc định là 48) — Số lượng lớp khối Transformer cần sử dụng.
  • activation_fn (str, mặc định là "gelu-approximate") — Hàm kích hoạt được sử dụng trong mạng truyền thẳng.
  • qk_norm (str, mặc định là "rms_norm_across_heads") — Lớp chuẩn hóa cần sử dụng.
  • rope_type (str, mặc định là "interleaved") — Phương thức áp dụng RoPE cần sử dụng ("interleaved" hoặc "split").

Transformer tinh chỉnh tự hồi quy theo chunk-causal được sử dụng làm giai đoạn 2 của SANA-WM.

Về mặt kiến trúc, hoàn toàn giống với LTX2VideoTransformer3DModel — cùng các đối số cấu hình, cùng các module con, cùng tên tham số — do đó một checkpoint LTX-2 đã phát hành có thể được nạp vào mà không cần thay đổi. Điều khác biệt nằm ở quá trình truyền tiến (forward pass):

  • chỉ luồng video được chạy (các nhánh chú ý chéo giữa âm thanh và âm thanh/video bị bỏ qua),
  • self-attention chạy trên một bộ nhớ đệm KV trượt cửa sổ rõ ràng (SanaWMRefinerKVCache) chứa attention sink cùng lịch sử tinh chỉnh gần nhất, do đó chi phí tinh chỉnh bị giới hạn cho mỗi khối AR và tăng tuyến tính theo độ dài video,
  • người gọi cung cấp video RoPE, cho phép mỗi cửa sổ AR giữ lại chỉ số tuyệt đối của từng khung hình trong video nguồn (xem build_rotary_emb_for_absolute_positions()).

Các đối số còn lại phản ánh LTX2VideoTransformer3DModel một cách tương ứng. Các đối số và mô-đun con phía âm thanh được giữ lại hoàn toàn để trọng số âm thanh của checkpoint có thể được lưu trữ và tải lại; chúng không được sử dụng bởi tiến trình forward của refiner.

build_rotary_emb_for_absolute_positions

< >

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

Tham số

  • batch_size (int) — Kích thước batch để phát sóng các tọa độ.
  • frame_positions (list[int]) — Chỉ số khung tiềm năng tuyệt đối được bao phủ bởi cửa sổ này.
  • height (int), width (int) — Độ phân giải không gian của latent.
  • device (torch.device) — Thiết bị để xây dựng các tọa độ trên đó.
  • fps (float) — Tốc độ khung hình của video, quyết định việc mở rộng RoPE theo thời gian của LTX-2.

Trả về

tuple[torch.Tensor, torch.Tensor]

bộ đôi RoPE (cos, sin).

Xây dựng video RoPE cho một danh sách rõ ràng các chỉ số khung tiềm năng tuyệt đối.

SanaWMLTX2AudioVideoRotaryPosEmbed.prepare_video_coords giả định một torch.arange(num_frames) liên tục, điều này là hợp lý cho suy luận hai chiều. Bộ tinh chỉnh AR cửa sổ trượt thay vào đó cần giữ lại chỉ số tuyệt đối của từng khung hình trong video nguồn, để RoPE nắm bắt đúng pha thời gian xuyên suốt cửa sổ [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

Tham số

  • hidden_states (torch.Tensor) — Latent video đã được patchify của khối đang hoạt động, có hình dạng (batch_size, num_video_tokens, in_channels).
  • encoder_hidden_states (torch.Tensor) — Nhúng văn bản có hình dạng (batch_size, text_seq_len, caption_channels).
  • timestep (torch.Tensor) — Timestep có shape (batch_size, num_video_tokens), đã được nhân với self.config.timestep_scale_multiplier.
  • video_rotary_emb (tuple[torch.Tensor, torch.Tensor]) — (cos, sin) RoPE cho các vị trí khung tuyệt đối của khối đang hoạt động, như được trả về bởi build_rotary_emb_for_absolute_positions().
  • encoder_attention_mask (torch.Tensor, tùy chọn) — Mặt nạ chú ý văn bản nhân có hình dạng (batch_size, text_seq_len).
  • kv_cache (SanaWMRefinerKVCache, optional) — Bộ nhớ đệm KV dạng cửa sổ trượt, lưu giữ attention sink và lịch sử tinh chỉnh gần nhất cho từng lớp, cùng với RoPE sink_pe dùng chung. Bắt buộc phải có khi kv_cache_mode được đặt.
  • kv_cache_mode (str, tuỳ chọn) — Một trong các giá trị sau:
    • "inject": chú ý đến K/V của [sink + history + current] (các bước khử nhiễu).
    • "capture_pre_rope": không có tiền tố; ghi lại K/V trước RoPE của bước forward này vào bộ nhớ đệm (được dùng một lần để khởi tạo attention sink từ các latents thô của giai đoạn 1).
    • "inject_and_capture_post_rope": chú ý đến [sink + history + current] K/V và ghi lại K/V sau RoPE của khối này vào cache để có thể thêm vào lịch sử.

    Khi None, khối thực hiện full self-attention thông thường chỉ trên các token hiện tại.

  • attention_kwargs (dict, tùy chọn) — Các kwargs tùy chọn được chuyển tiếp đến việc xử lý tỷ lệ LoRA.
  • return_dict (bool, tùy chọn, mặc định là True) — Có nên trả về một Transformer2DModelOutput thay vì một tuple thông thường hay không.

Trả về

Transformer2DModelOutput hoặc tuple

tốc độ dự đoán cho khối đang hoạt động, có hình dạng (batch_size, num_video_tokens, out_channels).

Đi qua tiến trình forward chỉ với video trên một khối AR đơn lẻ.

SanaWMRefinerKVCache

lớp diffusers.models.transformers.transformer_sana_wm_refiner.SanaWMRefinerKVCache

< >

( num_layers: int )

Tham số

  • num_layers (int) — Số lượng khối transformer để cấp phát bộ nhớ đệm cho từng lớp.

Container chứa một [SanaWMRefinerKVLayerCache] cho mỗi khối transformer, cùng với RoPE sink dùng chung.

Phần này triển khai hợp đồng KV-cache rf_shifted_sink mà bộ tinh chỉnh giai đoạn 2 của SANA-WM được huấn luyện với. Quá trình tinh chỉnh là nhân quả theo khối: block_size khung latent được khử nhiễu mỗi lần trong khi chú ý đến một cửa sổ có giới hạn gồm [attention sink + lịch sử gần đây + khối đang hoạt động] K/V.

  • sink_pe: bộ đôi (cos, sin) RoPE cho các khung sink, được xây dựng lại cho mỗi cửa sổ AR tại sink_rope_offset trượt để khung sink nằm ngay trước bộ nhớ đệm làm việc có giới hạn. Được dùng chung giữa các lớp vì RoPE không phụ thuộc vào lớp.

SanaWMRefinerKVLayerCache

lớp diffusers.models.transformers.transformer_sana_wm_refiner.SanaWMRefinerKVLayerCache

< >

( )

Bộ nhớ đệm KV theo từng lớp cho bộ tinh chỉnh AR nhân quả theo khối ở giai đoạn 2 của SANA-WM.

Chứa hai phần của tiền tố cửa sổ trượt mà self-attention của refiner chú ý đến, cùng một slot để đọc lại K/V mà lần forward cuối cùng đã ghi nhận. Tất cả các tensor đều có dạng (batch_size, num_tokens, inner_dim) (tức là trước khi unflatten theo từng head), khớp với bố cục mà self-attention của refiner nối vào.

  • sink_k_pre / sink_v: pre-RoPE K/V của các khung attention-sink, được ghi nhận một lần từ các latents thô của giai đoạn 1. Chúng được lưu ở dạng pre-RoPE để mỗi cửa sổ AR có thể áp dụng lại RoPE tại độ lệch sink đã dịch chuyển của riêng nó (SanaWMRefinerKVCache.sink_pe).
  • history_k / history_v: K/V post-RoPE của các khung gần đây đã được tinh chỉnh, sẵn sàng để nối trực tiếp.
  • captured_k_pre / captured_v_pre và captured_k_post / captured_v_post: các ô đọc lại được ghi bởi các kv_cache_mode bắt.

get_captured_post_rope

< >

( )

Lấy ra K/V sau RoPE được ghi nhận bởi bước forward cuối cùng.

get_captured_pre_rope

< >

( )

Lấy ra K/V trước RoPE được ghi nhận bởi bước forward cuối cùng.

get_history

< >

( )

Trả về K/V lịch sử gần đây sau RoPE, hoặc None nếu trống.

get_sink

< >

( )

Trả về K/V sink trước RoPE, hoặc None nếu nó chưa được ghi nhận (hoặc rỗng).

store_captured_post_rope

< >

( key: Tensorvalue: Tensor )

Lưu trữ K/V sau RoPE được tạo ra bởi bước forward hiện tại.

store_captured_pre_rope

< >

( key: Tensorvalue: Tensor )

Lưu trữ K/V trước RoPE được tạo ra bởi bước forward hiện tại.

store_history

< >

( history_k: Tensorhistory_v: Tensor )

Lưu trữ K/V lịch sử gần đây sau RoPE.

store_sink

< >

( sink_k_pre: Tensorsink_v: Tensor )

Lưu trữ K/V sink trước RoPE.

Transformer2DModelOutput

lớp diffusers.models.modeling_outputs.Transformer2DModelOutput

< >

( sample: torch.Tensor )

Tham số

  • sample (torch.Tensor có hình dạng (batch_size, num_channels, height, width) hoặc (batch size, num_vector_embeds - 1, num_latent_pixels) nếu Transformer2DModel là rời rạc) — Trạng thái ẩn đầu ra được điều kiện bởi đầu vào encoder_hidden_states. Nếu là rời rạc, trả về các phân phối xác suất cho các pixel tiềm ẩn không nhiễu.

Đầu ra của Transformer2DModel.

Update on GitHub

←SanaVideoTransformer3DModel SanaWMTransformer3DModel→