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

Le DiT autoregressif causal par blocs utilisé comme étape 2 de SanaWMPipeline, piloté par SanaWMLTX2Refiner.

Il est architecturalement identique à LTX2VideoTransformer3DModel — mêmes arguments de configuration, mêmes sous-modules, mêmes noms de paramètres — de sorte qu'un point de contrôle LTX-2 publié s'y charge sans modification. Le passage avant diffère :

  • seule la vidéo est traitée (les branches d'attention croisée audio et audio/vidéo sont ignorées),
  • l’auto-attention s’exécute contre un cache KV à fenêtre glissante explicite (SanaWMRefinerKVCache) contenant le puits d’attention ainsi que l’historique récent raffiné, de sorte que le calcul par bloc est borné et que le coût total de raffinement évolue linéairement avec la longueur de la vidéo,
  • l’appelant fournit le RoPE vidéo, ce qui permet à chaque fenêtre autorégressive de conserver l’index absolu de chaque image dans la vidéo source (voir SanaWMLTX2RefinerTransformer3DModel.build_rotary_emb_for_absolute_positions()).

Le modèle peut être chargé avec :

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 )

Paramètres

  • in_channels (int, par défaut 128) — Le nombre de canaux dans l'entrée.
  • out_channels (int, par défaut 128) — Le nombre de canaux dans la sortie.
  • patch_size (int, par défaut 1) — La taille des patches spatiaux à utiliser dans la couche d'encodage de patches.
  • patch_size_t (int, par défaut 1) — La taille des patches temporels à utiliser dans la couche d'embedding de patches.
  • num_attention_heads (int, par défaut 32) — Le nombre de têtes à utiliser pour l'attention multi-têtes.
  • attention_head_dim (int, par défaut 128) — Le nombre de canaux dans chaque tête.
  • cross_attention_dim (int, par défaut 4096) — Le nombre de canaux pour les têtes d'attention croisée.
  • num_layers (int, par défaut 48) — Le nombre de blocs Transformer à utiliser.
  • activation_fn (str, par défaut "gelu-approximate") — Fonction d'activation à utiliser dans le réseau de neurones à propagation avant.
  • qk_norm (str, par défaut "rms_norm_across_heads") — La couche de normalisation à utiliser.
  • rope_type (str, par défaut "interleaved") — Type d'application de RoPE à utiliser ("interleaved" ou "split").

Le transformateur autoregressif raffineur causal par blocs utilisé comme étape 2 de SANA-WM.

Architecturalement identique à LTX2VideoTransformer3DModel — mêmes arguments de configuration, mêmes sous-modules, mêmes noms de paramètres — de sorte qu'un point de contrôle LTX-2 publié s'y charge sans modification. Ce qui diffère, c'est le passage avant :

  • seule la vidéo est traitée (les branches d'attention croisée audio et audio/vidéo sont ignorées),
  • l’auto-attention s’exécute contre un cache KV à fenêtre glissante explicite (SanaWMRefinerKVCache) contenant le puits d’attention ainsi que l’historique récent raffiné, de sorte que le coût de raffinement est borné par bloc AR et évolue linéairement avec la longueur de la vidéo,
  • l’appelant fournit le RoPE vidéo, ce qui permet à chaque fenêtre AR de conserver l’index absolu de chaque image dans la vidéo source (voir build_rotary_emb_for_absolute_positions()).

Les arguments restants reflètent ceux de LTX2VideoTransformer3DModel un pour un. Les arguments et sous-modules du côté audio sont conservés uniquement pour que les poids audio du point de contrôle soient préservés lors de la sérialisation/désérialisation ; ils ne sont pas utilisés par le passage avant du raffineur.

build_rotary_emb_for_absolute_positions

< >

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

Paramètres

  • batch_size (int) — Taille du lot vers laquelle diffuser les coordonnées.
  • frame_positions (list[int]) — Indices absolus des trames latentes couverts par cette fenêtre.
  • height (int), width (int) — Résolution spatiale latente.
  • device (torch.device) — Appareil sur lequel construire les coordonnées.
  • fps (float) — Fréquence d’images de la vidéo, qui pilote le redimensionnement temporel du RoPE de LTX-2.

Retourne

tuple[torch.Tensor, torch.Tensor]

le tuple RoPE (cos, sin).

Construire le RoPE vidéo pour une liste explicite d’indices absolus de trames latentes.

SanaWMLTX2AudioVideoRotaryPosEmbed.prepare_video_coords suppose un torch.arange(num_frames) contigu, ce qui convient pour une inférence bidirectionnelle. Le raffineur AR à fenêtre glissante doit en revanche conserver l’index absolu de chaque image dans la vidéo source, afin que RoPE capture la phase temporelle correcte sur la fenêtre [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

Paramètres

  • hidden_states (torch.Tensor) — Latents vidéo patchifiés du bloc actif, de forme (batch_size, num_video_tokens, in_channels).
  • encoder_hidden_states (torch.Tensor) — Embeddings de texte de forme (batch_size, text_seq_len, caption_channels).
  • timestep (torch.Tensor) — Timestep de forme (batch_size, num_video_tokens), déjà mis à l'échelle par self.config.timestep_scale_multiplier.
  • video_rotary_emb (tuple[torch.Tensor, torch.Tensor]) — Le RoPE (cos, sin) pour les positions absolues du cadre actif, tel que renvoyé par build_rotary_emb_for_absolute_positions().
  • encoder_attention_mask (torch.Tensor, facultatif) — Masque d'attention textuel multiplicatif de forme (batch_size, text_seq_len).
  • kv_cache (SanaWMRefinerKVCache, facultatif) — Cache KV à fenêtre glissante contenant le puits d’attention par couche et l’historique récent raffiné, ainsi que le RoPE sink_pe partagé. Obligatoire dès que kv_cache_mode est défini.
  • kv_cache_mode (str, facultatif) — L'une des valeurs suivantes :
    • "inject" : attention sur [sink + history + current] K/V (les étapes de débruitage).
    • "capture_pre_rope" : aucun préfixe ; enregistre les K/V pré-RoPE de ce forward dans le cache (utilisé une fois pour initialiser le puits d’attention à partir des latents bruts de l’étape 1).
    • "inject_and_capture_post_rope" : porte attention sur [sink + history + current] K/V et enregistre les K/V post-RoPE de ce bloc dans le cache afin qu’ils puissent être ajoutés à l’historique.

    Lorsque None, le bloc exécute une auto-attention complète standard uniquement sur les jetons actuels.

  • attention_kwargs (dict, facultatif) — Arguments facultatifs transmis à la gestion de l'échelle LoRA.
  • return_dict (bool, facultatif, par défaut True) — Indique s'il faut retourner un Transformer2DModelOutput plutôt qu'un simple tuple.

Retourne

Transformer2DModelOutput ou tuple

la vitesse prédite pour le bloc actif, de forme (batch_size, num_video_tokens, out_channels).

Passage en avant vidéo uniquement sur un seul bloc AR.

SanaWMRefinerKVCache

classe diffusers.models.transformers.transformer_sana_wm_refiner.SanaWMRefinerKVCache

< >

( num_layers: int )

Paramètres

  • num_layers (int) — Nombre de blocs de transformeur pour lesquels allouer un cache par couche.

Conteneur contenant un [SanaWMRefinerKVLayerCache] par bloc de transformeur, ainsi que le RoPE de sink partagé.

Ceci implémente le contrat de cache KV rf_shifted_sink avec lequel le raffineur de l’étape 2 de SANA-WM a été entraîné. Le raffinement est causal par blocs : block_size trames latentes sont débruitées à la fois en tenant compte d’une fenêtre bornée de [sink d’attention + historique récent + bloc actif] K/V.

  • sink_pe : le tuple RoPE (cos, sin) pour les frames de sink, reconstruit à chaque fenêtre AR au décalage glissant sink_rope_offset afin que le sink se situe immédiatement avant le cache de travail borné. Partagé entre les couches car RoPE ne dépend pas de la couche.

SanaWMRefinerKVLayerCache

classe diffusers.models.transformers.transformer_sana_wm_refiner.SanaWMRefinerKVLayerCache

< >

( )

Cache KV par couche pour le raffineur AR causal par blocs de l'étape 2 de SANA-WM.

Contient les deux moitiés du préfixe de fenêtre glissante auquel l'auto-attention du refiner se réfère, ainsi qu'un emplacement pour relire les K/V capturés par le dernier forward. Tous les tenseurs sont de forme (batch_size, num_tokens, inner_dim) (c'est-à-dire avant le déplateau par tête), correspondant à la disposition que l'auto-attention du refiner concatène.

  • sink_k_pre / sink_v : pre-RoPE K/V des frames de sink d’attention, capturées une seule fois à partir des latents bruts de l’étape 1. Elles sont stockées avant application de RoPE afin que chaque fenêtre AR puisse réappliquer RoPE à son propre décalage de sink (SanaWMRefinerKVCache.sink_pe).
  • history_k / history_v : K/V post-RoPE des frames récentes déjà affinées, prêtes à être concaténées telles quelles.
  • captured_k_pre / captured_v_pre et captured_k_post / captured_v_post : emplacements de lecture écrits par les modes de capture kv_cache_mode.

get_captured_post_rope

< >

( )

Extrait les K/V post-RoPE capturés par le dernier forward.

get_captured_pre_rope

< >

( )

Extrait les K/V pré-RoPE capturés par le dernier passage en avant.

get_history

< >

( )

Retourner l’historique récent K/V post-RoPE, ou None s’il est vide.

get_sink

< >

( )

Retourne le K/V de sink pré-RoPE, ou None s'il n'a pas été capturé (ou s'il est vide).

store_captured_post_rope

< >

( key: Tensorvalue: Tensor )

Stockez les K/V post-RoPE produits par le forward actuel.

store_captured_pre_rope

< >

( key: Tensorvalue: Tensor )

Stockez les K/V pré-RoPE produits par le forward actuel.

store_history

< >

( history_k: Tensorhistory_v: Tensor )

Stockez l’historique récent des K/V après application de RoPE.

store_sink

< >

( sink_k_pre: Tensorsink_v: Tensor )

Stockez le K/V de sink pré-RoPE.

Transformer2DModelOutput

classe diffusers.models.modeling_outputs.Transformer2DModelOutput

< >

( sample: torch.Tensor )

Paramètres

  • sample (torch.Tensor de forme (batch_size, num_channels, height, width) ou (batch size, num_vector_embeds - 1, num_latent_pixels) si Transformer2DModel est discret) — Les états cachés en sortie, conditionnés par l'entrée encoder_hidden_states. Si discret, renvoie les distributions de probabilité pour les pixels latents non bruités.

La sortie de Transformer2DModel.

Update on GitHub

←SanaVideoTransformer3DModel SanaWMTransformer3DModel→