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

El refiner DiT autoregresivo causal por bloques utilizado como etapa 2 de SanaWMPipeline, impulsado por SanaWMLTX2Refiner.

Es arquitectónicamente idéntico a LTX2VideoTransformer3DModel: mismos argumentos de configuración, mismos submódulos, mismos nombres de parámetros, por lo que un checkpoint de LTX-2 publicado se carga en él sin cambios. El paso hacia adelante (forward pass) difiere:

  • solo se ejecuta el flujo de video (se omiten las ramas de atención cruzada de audio y audio/video),
  • self-attention se ejecuta contra una caché KV de ventana deslizante explícita (SanaWMRefinerKVCache) que mantiene el sink de atención junto con el historial refinado reciente, de modo que el cómputo por bloque está acotado y el coste total de refinamiento escala linealmente con la longitud del video,
  • el llamador proporciona el RoPE del video, lo que permite que cada ventana autoregresiva conserve el índice absoluto de cada fotograma en el video de origen (ver SanaWMLTX2RefinerTransformer3DModel.build_rotary_emb_for_absolute_positions()).

El modelo se puede cargar con:

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

clase 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, por defecto 128) — El número de canales en la entrada.
  • out_channels (int, por defecto 128) — El número de canales en la salida.
  • patch_size (int, por defecto 1) — El tamaño de los parches espaciales que se usarán en la capa de incrustación de parches.
  • patch_size_t (int, por defecto 1) — El tamaño de los parches temporales a utilizar en la capa de incrustación de parches.
  • num_attention_heads (int, por defecto 32) — El número de cabezas a utilizar para la atención multi-cabeza.
  • attention_head_dim (int, por defecto 128) — El número de canales en cada cabeza.
  • cross_attention_dim (int, por defecto 4096) — El número de canales para las cabezas de atención cruzada.
  • num_layers (int, por defecto 48) — El número de capas de bloques Transformer a utilizar.
  • activation_fn (str, por defecto "gelu-approximate") — Función de activación a utilizar en la red feed-forward.
  • qk_norm (str, por defecto "rms_norm_across_heads") — La capa de normalización a utilizar.
  • rope_type (str, por defecto "interleaved") — Qué aplicación de RoPE utilizar ("interleaved" o "split").

El transformer autoregresivo de refinamiento con causalidad por bloques utilizado como etapa 2 de SANA-WM.

Arquitectónicamente idéntico a LTX2VideoTransformer3DModel: mismos argumentos de configuración, mismos submódulos, mismos nombres de parámetros, por lo que un checkpoint de LTX-2 publicado se carga en él sin cambios. Lo que difiere es la pasada hacia adelante:

  • solo se ejecuta el flujo de video (se omiten las ramas de atención cruzada de audio y audio/video),
  • self-attention se ejecuta contra una caché KV de ventana deslizante explícita (SanaWMRefinerKVCache) que mantiene el sumidero de atención junto con el historial refinado reciente, de modo que el costo de refinamiento está acotado por bloque AR y escala linealmente con la longitud del video,
  • el llamador proporciona el RoPE del video, lo que permite que cada ventana de AR conserve el índice absoluto de cada cuadro en el video de origen (consulte build_rotary_emb_for_absolute_positions()).

Los argumentos restantes replican uno a uno los de LTX2VideoTransformer3DModel. Los argumentos y submódulos del lado de audio se conservan únicamente para que los pesos de audio del checkpoint se conserven al serializar y deserializar; no se utilizan en la propagación hacia adelante del 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) — Tamaño del lote al que se deben difundir las coordenadas.
  • frame_positions (list[int]) — Índices absolutos de fotogramas latentes cubiertos por esta ventana.
  • height (int), width (int) — Resolución espacial latente.
  • device (torch.device) — Dispositivo en el que se construirán las coordenadas.
  • fps (float) — Frecuencia de fotogramas del vídeo, que determina la escalada temporal de RoPE en LTX-2.

Devuelve

tuple[torch.Tensor, torch.Tensor]

la tupla RoPE (cos, sin).

Construya el RoPE de video para una lista explícita de índices absolutos de fotogramas latentes.

SanaWMLTX2AudioVideoRotaryPosEmbed.prepare_video_coords asume un torch.arange(num_frames) contiguo, lo cual es adecuado para la inferencia bidireccional. El refinador AR de ventana deslizante, en cambio, necesita mantener el índice absoluto de cada fotograma en el video de origen, de modo que RoPE capture la fase temporal correcta a lo largo de la ventana [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 video parcheados del bloque activo, de forma (batch_size, num_video_tokens, in_channels).
  • encoder_hidden_states (torch.Tensor) — Incrustaciones de texto de forma (batch_size, text_seq_len, caption_channels).
  • timestep (torch.Tensor) — Paso temporal de forma (batch_size, num_video_tokens), ya escalado por self.config.timestep_scale_multiplier.
  • video_rotary_emb (tuple[torch.Tensor, torch.Tensor]) — El RoPE (cos, sin) para las posiciones absolutas del marco activo, tal como lo devuelve build_rotary_emb_for_absolute_positions().
  • encoder_attention_mask (torch.Tensor, opcional) — Máscara de atención multiplicativa para el texto, de forma (batch_size, text_seq_len).
  • kv_cache (SanaWMRefinerKVCache, opcional) — Cache KV de ventana deslizante que almacena el attention sink por capa y el historial refinado reciente, además del RoPE compartido sink_pe. Obligatorio siempre que se establezca kv_cache_mode.
  • kv_cache_mode (str, opcional) — Uno de:
    • "inject": prestar atención a [sink + history + current] K/V (los pasos de desruido).
    • "capture_pre_rope": sin prefijo; registra los K/V pre-RoPE de esta pasada hacia adelante en la caché (se usa una vez para sembrar el sumidero de atención a partir de los latentes crudos de la etapa 1).
    • "inject_and_capture_post_rope": atiende a [sink + history + current] K/V y registra los K/V post-RoPE de este bloque en la caché para que puedan agregarse al historial.

    Cuando es None, el bloque ejecuta una autoatención completa simple solo sobre los tokens actuales.

  • attention_kwargs (dict, opcional) — Kwargs opcionales que se reenvían al manejo de la escala de LoRA.
  • return_dict (bool, opcional, valor predeterminado True) — Si se debe devolver un Transformer2DModelOutput en lugar de una tupla simple.

Devuelve

Transformer2DModelOutput o tuple

la velocidad predicha para el bloque activo, de forma (batch_size, num_video_tokens, out_channels).

Paso de propagación hacia adelante solo de video sobre un solo bloque AR.

SanaWMRefinerKVCache

clase diffusers.models.transformers.transformer_sana_wm_refiner.SanaWMRefinerKVCache

< >

( num_layers: int )

Parámetros

  • num_layers (int) — Número de bloques de transformador para los que se asignará una caché por capa.

Contenedor que mantiene una instancia de SanaWMRefinerKVLayerCache por bloque de transformador, además del RoPE de sink compartido.

Esto implementa el contrato de KV-cache rf_shifted_sink con el que se entrenó el refinador de la etapa 2 de SANA-WM. El refinamiento es causal por bloques: block_size marcos latentes se desruidan a la vez mientras atienden a una ventana acotada de [attention sink + historial reciente + bloque activo] K/V.

  • sink_pe: la tupla RoPE (cos, sin) para los cuadros de sink, reconstruida en cada ventana AR en el desplazamiento deslizante sink_rope_offset para que el sink se sitúe inmediatamente antes de la caché de trabajo acotada. Compartida entre capas porque RoPE no depende de la capa.

SanaWMRefinerKVLayerCache

clase diffusers.models.transformers.transformer_sana_wm_refiner.SanaWMRefinerKVLayerCache

< >

( )

Cache KV por capa para el refinador AR causal por bloques de la etapa 2 de SANA-WM.

Contiene las dos mitades del prefijo de ventana deslizante al que atiende la autoatención del refiner, además de un espacio para leer de nuevo los K/V que capturó la última forward. Todos los tensores son (batch_size, num_tokens, inner_dim) (es decir, antes del unflatten por cabeza), coincidiendo con el layout que la autoatención del refiner concatena.

  • sink_k_pre / sink_v: pre-RoPE K/V de los fotogramas de attention-sink, capturados una sola vez desde los latentes crudos de la etapa 1. Se almacenan pre-RoPE para que cada ventana AR pueda volver a aplicar RoPE en su propio desplazamiento de offset de sink (SanaWMRefinerKVCache.sink_pe).
  • history_k / history_v: K/V post-RoPE de los fotogramas recientes ya refinados, listos para concatenarse tal cual.
  • captured_k_pre / captured_v_pre y captured_k_post / captured_v_post: ranuras de lectura escritas por los modos de captura kv_cache_mode.

get_captured_post_rope

< >

( )

Extrae los K/V posteriores a RoPE capturados por la última pasada hacia adelante.

get_captured_pre_rope

< >

( )

Extrae los K/V pre-RoPE capturados por el último paso de adelante.

get_history

< >

( )

Devuelve el K/V de historial reciente posterior a RoPE, o None si está vacío.

get_sink

< >

( )

Devuelve el K/V de sink pre-RoPE, o None si no se ha capturado (o está vacío).

store_captured_post_rope

< >

( key: Tensorvalue: Tensor )

Almacene el K/V posterior a RoPE generado por la propagación hacia adelante actual.

store_captured_pre_rope

< >

( key: Tensorvalue: Tensor )

Almacene los K/V pre-RoPE generados por la propagación hacia adelante actual.

store_history

< >

( history_k: Tensorhistory_v: Tensor )

Almacene el K/V del historial reciente posterior a RoPE.

store_sink

< >

( sink_k_pre: Tensorsink_v: Tensor )

Almacene el K/V del sink pre-RoPE.

Transformer2DModelOutput

clase diffusers.models.modeling_outputs.Transformer2DModelOutput

< >

( sample: torch.Tensor )

Parámetros

  • sample (torch.Tensor de forma (batch_size, num_channels, height, width) o (batch size, num_vector_embeds - 1, num_latent_pixels) si Transformer2DModel es discreto) — La salida de estados ocultos condicionada a la entrada encoder_hidden_states. Si es discreto, devuelve distribuciones de probabilidad para los píxeles latentes sin ruido.

La salida de Transformer2DModel.

Update on GitHub

←SanaVideoTransformer3DModel SanaWMTransformer3DModel→