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

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

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: list[int]height: 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: 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

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→