Hugging Face's logo

Diffusers documentation

SanaWMTransformer3DModel

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

SanaWMTransformer3DModel

Un Diffusion Transformer 3D (1,6 Md de paramètres) pour la génération d'images en vidéo contrôlée par caméra, utilisé comme échantillonneur de l'étape 1 de SanaWMPipeline. Le transformer combine :

  • une branche principale d'attention linéaire GDN-Triton bidirectionnelle (profondeur 20, 2240 cachés, 20 têtes),
  • une branche de contrôle de caméra UCPE (Unified Camera Pose Embedding) qui consomme une représentation raymap + Plücker de la trajectoire demandée, et
  • un embedding de position rotatif 3D de type Wan, ainsi que des blocs d'attention softmax périodiques injectés toutes les softmax_every_n couches.

Le modèle peut être chargé avec :

import torch
from diffusers import SanaWMTransformer3DModel

transformer = SanaWMTransformer3DModel.from_pretrained(
    "Efficient-Large-Model/SANA-WM_bidirectional-diffusers",
    subfolder="transformer",
    torch_dtype=torch.bfloat16,
)

SanaWMTransformer3DModel

classe diffusers.SanaWMTransformer3DModel

< >

( in_channels: int = 128num_layers: int = 20hidden_size: int = 2240num_attention_heads: int = 20patch_size: tuple[int, int, int] = (1, 1, 1)softmax_every_n: int = 4linear_head_dim: int = 112ffn_type: str = 'GLUMBConvTemp't_kernel_size: int = 3conv_kernel_size: int = 4k_conv_only: bool = Truepos_embed_type: str = 'wan_rope'qk_norm: bool = Truecross_norm: bool = Truey_norm: bool = Truecam_attn_compress: int = 1init_cam_from_base: bool = Trueuse_chunk_plucker_post_attn: bool = Truechunk_plucker_channels: int = 48chunk_plucker_post_attn_blocks: int = 20fp32_attention: bool = Trueimage_size: int = 720caption_channels: int = 2304mlp_ratio: float = 3.0mlp_acts: tuple = ('silu', 'silu', None)use_pe: bool = Truelearn_sigma: bool = Falsepred_sigma: bool = Falsemixed_precision: str = 'bf16' )

Paramètres

  • in_channels (int, par défaut 128) — Canaux latents du VAE (LTX-2).
  • softmax_every_n (int, par défaut 4) — Utilise un bloc d'attention softmax tous les N blocs.
  • linear_head_dim (int, par défaut 112) — Dimension des têtes GDN.
  • ffn_type (str, par défaut "GLUMBConvTemp") — FFN.
  • t_kernel_size (int, par défaut 3) — Noyau de convolution temporelle.
  • conv_kernel_size (int, par défaut 4) — Taille du noyau de convolution spatiale dans l'attention.
  • k_conv_only (bool, par défaut True) — Appliquer la convolution uniquement sur K.
  • pos_embed_type (str, par défaut "wan_rope") — Incrustation de position.
  • qk_norm (bool, par défaut True) — RMSNorm sur Q/K.
  • cross_norm (bool, par défaut True) — RMSNorm sur K de l'attention croisée.
  • y_norm (bool, par défaut True) — Applique attention_y_norm aux embeddings de texte.
  • init_cam_from_base (bool, par défaut True) — Non utilisé ; la branche caméra est chargée depuis le point de contrôle. Conservé afin que les fichiers config.json publiés puissent être chargés.
  • use_chunk_plucker_post_attn (bool, par défaut True). —
  • chunk_plucker_channels (int, par défaut 48) — 6 dims * temporal_stride 8.
  • chunk_plucker_post_attn_blocks (int, par défaut 20) — Tous les blocs.
  • fp32_attention (bool, par défaut True) — Non utilisé ; l'attention s'exécute toujours en fp32. Conservé afin que les fichiers config.json publiés puissent être chargés.
  • image_size (int, par défaut 720) — Taille nominale de l'image.
  • caption_channels (int, par défaut 2304) — Taille cachée de Gemma-2.

SANA-WM 1600M DiT bidirectionnel contrôlé par caméra.

Un DiT à classe unique (depth=20, hidden_size=2240, patch_size=(1,1,1), num_heads=20 — c’est-à-dire la version publique Efficient-Large-Model/SANA-WM_bidirectional). save_pretrained / from_pretrained fonctionnent immédiatement grâce à ConfigMixin.

Chaque bloc exécute une attention conditionnée par la caméra (UCPE) : [BidirectionalGDNUCPESinglePathLiteLA], à l'exception de chaque softmax_every_n-ième bloc, qui exécute son équivalent softmax [_SoftmaxUCPESinglePathLiteLA].

Le state-dict est identique au checkpoint public sana, à l'exception du tampon pos_embed supprimé intentionnellement.

forward

< >

( hidden_states: torch.Tensortimestep: torch.Tensorencoder_hidden_states: torch.Tensorencoder_attention_mask: torch.Tensor | None = Nonemask: torch.Tensor | None = Nonereturn_dict: bool = Truedata_info: Optional[dict] = Nonecamera_conditions: torch.Tensor | None = Nonechunk_plucker: torch.Tensor | None = Nonecam_pos_embeds: Optional[dict] = Nonepos_embeds: torch.Tensor | None = Noneraymats: torch.Tensor | None = Noneframe_valid_mask: torch.Tensor | None = Nonechunk_size: int | None = None )

Paramètres

  • hidden_states — latents de forme (B, C, T, H, W).
  • timestep — (B, 1, T) timesteps de diffusion par image (style LTX).
  • encoder_hidden_states — embeddings de texte (B, 1, L, D_caption).
  • encoder_attention_mask — masque d'attention du texte (B, L) (convention diffusers).
  • mask — Alias de encoder_attention_mask correspondant au nom du paramètre du DiT sana. Si les deux sont fournis, mask a la priorité.
  • return_dict — Si True (par défaut), renvoie un Transformer2DModelOutput ; sinon, renvoie un tuple unique (sample,).
  • data_info — Conditionnement supplémentaire ; data_info["image_vae_embeds"] est concatené aux latents le long de l’axe des canaux lorsqu’il est présent.
  • camera_conditions — (B, T, 20) conditions brutes de la caméra pilotant la branche de contrôle de la caméra (UCPE).
  • chunk_plucker — embeddings de rayons Plucker (B, C, T, H, W), utilisés lorsque le modèle est configuré avec use_chunk_plucker_input / use_chunk_plucker_post_attn.
  • cam_pos_embeds — Incrustations positionnelles de caméra précalculées facultatives (entrées "absmap" / "P") réutilisées à la place d'un recalcul.
  • pos_embeds — Emplacements de position rotatifs précalculés facultatifs ; lorsque la valeur est None, ils sont construits à partir de la forme des latents.
  • raymats — Matrices de rayons UCPE précalculées (utilisées uniquement lorsque cam_pos_embeds ne contient pas "P").
  • frame_valid_mask — Masque de validité par image facultatif, de forme (B, 1, T, 1, 1), (B, 1, T) ou (B, T).
  • chunk_size — Remplacement de la longueur des blocs transmis aux blocs ; en cas d’absence, la valeur de chunk_size configurée pour chaque bloc est utilisée.

Exécuter le DiT de SANA-WM.

unpatchify

< >

( xpost_patch_num_framespost_patch_heightpost_patch_width )

x : (N, T, patch_size*2 C) imgs : (N, H, W, C)

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

←SanaWMLTX2RefinerTransformer3DModel SD3Transformer2DModel→