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.42.1).
Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

SanaWMLTX2RefinerTransformer3DModel

O refiner DiT autoregressivo causal por blocos usado como etapa 2 do SanaWMPipeline, dirigido por SanaWMLTX2Refiner.

Ele é arquitetonicamente idêntico ao LTX2VideoTransformer3DModel — mesmos argumentos de configuração, mesmos submódulos, mesmos nomes de parâmetros — de modo que um checkpoint LTX-2 publicado é carregado nele sem alterações. A passagem para frente difere:

  • somente o fluxo de vídeo é executado (os ramos de atenção cruzada de áudio e áudio/vídeo são ignorados),
  • o self-attention é executado contra um cache KV de janela deslizante explícito (SanaWMRefinerKVCache) que armazena o sink de atenção e o histórico refinado recente, de modo que o cálculo por bloco é limitado e o custo total de refinamento escala linearmente com o comprimento do vídeo,
  • o chamador fornece o RoPE do vídeo, o que permite que cada janela autoregressiva mantenha o índice absoluto de cada quadro no vídeo de origem (veja SanaWMLTX2RefinerTransformer3DModel.build_rotary_emb_for_absolute_positions()).

O modelo pode ser carregado com:

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

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

Parâmetros

  • in_channels (int, padrão 128) — O número de canais na entrada.
  • out_channels (int, padrão 128) — O número de canais na saída.
  • patch_size (int, padrão 1) — O tamanho dos patches espaciais a serem usados na camada de incorporação de patches.
  • patch_size_t (int, padrão 1) — O tamanho dos patches temporais a serem usados na camada de incorporação de patches.
  • num_attention_heads (int, padrão 32) — O número de cabeças a serem usadas para a atenção multi-cabeça.
  • attention_head_dim (int, padrão 128) — O número de canais em cada head.
  • cross_attention_dim (int, padrão 4096) — O número de canais para as cabeças de atenção cruzada.
  • num_layers (int, padrão 48) — O número de blocos de camadas de Transformer a serem usados.
  • activation_fn (str, padrão "gelu-approximate") — Função de ativação a ser usada no feed-forward.
  • qk_norm (str, padrão "rms_norm_across_heads") — A camada de normalização a ser usada.
  • rope_type (str, padrão "interleaved") — Qual aplicação de RoPE usar ("interleaved" ou "split").

O transformer autoregressivo refinador causal por blocos usado como etapa 2 do SANA-WM.

Arquiteturalmente idêntico a LTX2VideoTransformer3DModel — mesmos argumentos de configuração, mesmos submódulos, mesmos nomes de parâmetros — de modo que um checkpoint LTX-2 publicado é carregado nele sem alterações. O que difere é o forward pass:

  • somente o fluxo de vídeo é executado (os ramos de atenção cruzada de áudio e áudio/vídeo são ignorados),
  • o self-attention é executado contra um cache KV de janela deslizante explícito (SanaWMRefinerKVCache) que armazena o sink de atenção e o histórico refinado recente, de modo que o custo de refinamento é limitado por bloco AR e escala linearmente com o comprimento do vídeo,
  • o chamador fornece o RoPE do vídeo, o que permite que cada janela AR mantenha o índice absoluto de cada quadro no vídeo de origem (veja build_rotary_emb_for_absolute_positions()).

Os argumentos restantes espelham os de LTX2VideoTransformer3DModel um a um. Os argumentos e submódulos do lado de áudio são mantidos puramente para que os pesos de áudio do checkpoint façam ida e volta; eles não são usados pelo forward do refiner.

build_rotary_emb_for_absolute_positions

< >

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

Parâmetros

  • batch_size (int) — Tamanho do lote para o qual as coordenadas devem ser transmitidas.
  • frame_positions (list[int]) — Índices absolutos de quadros latentes cobertos por esta janela.
  • height (int), width (int) — Resolução espacial latente.
  • device (torch.device) — Dispositivo em que as coordenadas serão construídas.
  • fps (float) — Taxa de quadros do vídeo, que controla a escalação do RoPE temporal do LTX-2.

Retorna

tuple[torch.Tensor, torch.Tensor]

a tupla RoPE (cos, sin).

Construa o RoPE de vídeo para uma lista explícita de índices absolutos de quadros latentes.

SanaWMLTX2AudioVideoRotaryPosEmbed.prepare_video_coords pressupõe um torch.arange(num_frames) contíguo, o que é adequado para inferência bidirecional. O refinador AR de janela deslizante, por sua vez, precisa manter o índice absoluto de cada quadro no vídeo de origem, para que o RoPE capture a fase temporal correta ao longo da janela [sink + recent + active].

forward

< >

( hidden_states: Tensorencoder_hidden_states: Tensortimestep: Tensorvideo_rotary_emb: tupleencoder_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

Parâmetros

  • hidden_states (torch.Tensor) — Latentes de vídeo em patches do bloco ativo, com forma (batch_size, num_video_tokens, in_channels).
  • encoder_hidden_states (torch.Tensor) — Incorporações de texto com forma (batch_size, text_seq_len, caption_channels).
  • timestep (torch.Tensor) — Timestep de forma (batch_size, num_video_tokens), já escalado por self.config.timestep_scale_multiplier.
  • video_rotary_emb (tuple[torch.Tensor, torch.Tensor]) — O RoPE (cos, sin) para as posições absolutas do quadro ativo, conforme retornado por build_rotary_emb_for_absolute_positions().
  • encoder_attention_mask (torch.Tensor, opcional) — Máscara de atenção multiplicativa do texto com forma (batch_size, text_seq_len).
  • kv_cache (SanaWMRefinerKVCache, opcional) — Cache KV de janela deslizante que armazena o attention sink por camada e o histórico refinado recente, além do sink_pe RoPE compartilhado. Obrigatório sempre que kv_cache_mode estiver definido.
  • kv_cache_mode (str, opcional) — Um dos seguintes:
    • "inject": atende a [sink + history + current] K/V (as etapas de desruido).
    • "capture_pre_rope": sem prefixo; registra o K/V pré-RoPE deste forward no cache (usado uma vez para semeiar o attention sink a partir dos latentes brutos do estágio 1).
    • "inject_and_capture_post_rope": atenda a [sink + history + current] K/V e registre o K/V pós-RoPE deste bloco no cache para que ele possa ser adicionado ao histórico.

    Quando None, o bloco executa uma autoatenção completa simples apenas sobre os tokens atuais.

  • attention_kwargs (dict, opcional) — Kwargs opcionais encaminhados para o tratamento da escala do LoRA.
  • return_dict (bool, opcional, padrão True) — Se deve retornar um Transformer2DModelOutput em vez de uma tupla simples.

Retorna

Transformer2DModelOutput ou tuple

a velocidade prevista para o bloco ativo, de forma (batch_size, num_video_tokens, out_channels).

Passada de propagação direta apenas de vídeo sobre um único bloco AR.

SanaWMRefinerKVCache

classe diffusers.models.transformers.transformer_sana_wm_refiner.SanaWMRefinerKVCache

< >

( num_layers: int )

Parâmetros

  • num_layers (int) — Número de blocos de transformador para os quais alocar um cache por camada.

Contêiner que mantém uma [SanaWMRefinerKVLayerCache] por bloco do transformer, além do RoPE de sink compartilhado.

Isso implementa o contrato de KV-cache rf_shifted_sink com o qual o refinador do estágio 2 do SANA-WM foi treinado. O refinamento é causal por blocos: block_size quadros latentes são desruidos de cada vez, enquanto prestam atenção a uma janela limitada de [attention sink + histórico recente + bloco ativo] K/V.

  • sink_pe: a tupla RoPE (cos, sin) para os quadros de sink, reconstruída a cada janela de AR no sink_rope_offset deslizante, de modo que o sink fique imediatamente antes do cache de trabalho limitado. Compartilhada entre as camadas, pois o RoPE não depende da camada.

SanaWMRefinerKVLayerCache

classe diffusers.models.transformers.transformer_sana_wm_refiner.SanaWMRefinerKVLayerCache

< >

( )

Cache de KV por camada para o refinador AR causal por blocos do estágio 2 do SANA-WM.

Contém as duas metades do prefixo da janela deslizante ao qual a autoatenção do refiner presta atenção, além de um slot para ler de volta o K/V capturado pelo último forward. Todos os tensors são (batch_size, num_tokens, inner_dim) (ou seja, antes do unflatten por cabeça), correspondendo ao layout que a autoatenção do refiner concatena.

  • sink_k_pre / sink_v: pre-RoPE K/V dos quadros de attention-sink, capturados uma única vez dos latentes brutos da etapa 1. Eles são armazenados pré-RoPE para que cada janela AR possa reaplicar o RoPE em seu próprio deslocamento de sink (SanaWMRefinerKVCache.sink_pe).
  • history_k / history_v: K/V post-RoPE dos quadros recentes já refinados, prontos para serem concatenados diretamente.
  • captured_k_pre / captured_v_pre e captured_k_post / captured_v_post: slots de leitura gravados pelos modos de captura kv_cache_mode.

get_captured_post_rope

< >

( )

Extrair o K/V pós-RoPE capturado pelo último forward.

get_captured_pre_rope

< >

( )

Pegue o K/V pré-RoPE capturado pelo último forward.

get_history

< >

( )

Retorna o K/V do histórico recente pós-RoPE, ou None se estiver vazio.

get_sink

< >

( )

Retorna o K/V de sink pré-RoPE, ou None se ele não tiver sido capturado (ou estiver vazio).

store_captured_post_rope

< >

( key: Tensorvalue: Tensor )

Armazene o K/V pós-RoPE produzido pelo forward atual.

store_captured_pre_rope

< >

( key: Tensorvalue: Tensor )

Armazene o K/V pré-RoPE produzido pelo forward atual.

store_history

< >

( history_k: Tensorhistory_v: Tensor )

Armazene o K/V do histórico recente pós-RoPE.

store_sink

< >

( sink_k_pre: Tensorsink_v: Tensor )

Armazene o K/V do sink pré-RoPE.

Transformer2DModelOutput

classe diffusers.models.modeling_outputs.Transformer2DModelOutput

< >

( sample: torch.Tensor )

Parâmetros

  • sample (torch.Tensor de forma (batch_size, num_channels, height, width) ou (batch size, num_vector_embeds - 1, num_latent_pixels) se Transformer2DModel for discreto) — A saída dos estados ocultos condicionada à entrada encoder_hidden_states. Se discreto, retorna distribuições de probabilidade para os pixels latentes sem ruído.

A saída de Transformer2DModel.

Update on GitHub

←SanaVideoTransformer3DModel SanaWMTransformer3DModel→