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

Um 3D Diffusion Transformer (1,6B parâmetros) para geração de imagem-para-vídeo controlada por câmera, usado como o amostrador da etapa 1 do SanaWMPipeline. O transformer combina:

  • um ramo principal de atenção linear bidirecional GDN-Triton (profundidade 20, oculto 2240, 20 cabeças),
  • a ramificação de controle de câmera UCPE (Unified Camera Pose Embedding) que consome uma representação raymap + Plücker da trajetória solicitada, e
  • um incorporação de posição rotacional 3D no estilo Wan, além de blocos de atenção softmax periódicos injetados a cada softmax_every_n camadas.

O modelo pode ser carregado com:

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

Parâmetros

  • in_channels (int, padrão 128) — Canais latentes do VAE (LTX-2).
  • softmax_every_n (int, padrão 4) — Use um bloco de atenção softmax a cada N blocos.
  • linear_head_dim (int, padrão 112) — Dimensão da cabeça do GDN.
  • ffn_type (str, padrão "GLUMBConvTemp") — FFN.
  • t_kernel_size (int, padrão 3) — Kernel da convolução temporal.
  • conv_kernel_size (int, padrão 4) — Kernel de convolução espacial dentro da atenção.
  • k_conv_only (bool, padrão True) — Aplique conv apenas em K.
  • pos_embed_type (str, padrão "wan_rope") — Incorporação de posição.
  • qk_norm (bool, padrão True) — RMSNorm em Q/K.
  • cross_norm (bool, padrão True) — RMSNorm no K da atenção cruzada.
  • y_norm (bool, padrão True) — Aplica attention_y_norm aos embeddings de texto.
  • init_cam_from_base (bool, padrão True) — Não utilizado; o ramo da câmera é carregado do checkpoint. Mantido para que os arquivos config.json liberados sejam carregados.
  • use_chunk_plucker_post_attn (bool, padrão True). —
  • chunk_plucker_channels (int, padrão 48) — 6 dims * temporal_stride 8.
  • chunk_plucker_post_attn_blocks (int, padrão 20) — Todos os blocos.
  • fp32_attention (bool, padrão True) — Não utilizado; a atenção sempre é executada em fp32. Mantido para que os arquivos config.json liberados sejam carregados.
  • image_size (int, padrão 720) — Tamanho nominal da imagem.
  • caption_channels (int, padrão 2304) — Tamanho oculto do Gemma-2.

SANA-WM 1600M DiT bidirecional controlado por câmera.

Um DiT de classe única (depth=20, hidden_size=2240, patch_size=(1,1,1), num_heads=20 — ou seja, a liberação pública Efficient-Large-Model/SANA-WM_bidirectional). save_pretrained / from_pretrained funcionam imediatamente via ConfigMixin.

Cada bloco executa uma atenção condicionada por câmera (UCPE): [BidirectionalGDNUCPESinglePathLiteLA], exceto a cada softmax_every_n-ésimo bloco, que executa seu equivalente de softmax [_SoftmaxUCPESinglePathLiteLA].

O state-dict é idêntico ao checkpoint público do sana, exceto pelo buffer pos_embed removido intencionalmente.

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 )

Parâmetros

  • hidden_states — latentes (B, C, T, H, W).
  • timestep — (B, 1, T) timesteps de difusão por quadro (estilo LTX).
  • encoder_hidden_states — embeddings de texto (B, 1, L, D_caption).
  • encoder_attention_mask — máscara de atenção do texto (B, L) (convenção do diffusers).
  • mask — Alias para encoder_attention_mask correspondente ao nome do kwarg do DiT do sana. Se ambos forem passados, mask tem precedência.
  • return_dict — Se True (padrão), retorna um Transformer2DModelOutput; caso contrário, retorna uma tupla única (sample,).
  • data_info — Condição extra; data_info["image_vae_embeds"] é concatenado aos latentes ao longo do eixo de canais quando presente.
  • camera_conditions — (B, T, 20) condições brutas da câmera que controlam o ramo de controle de câmera (UCPE).
  • chunk_plucker — embeddings de raios Plucker (B, C, T, H, W), consumidos quando o modelo é configurado com use_chunk_plucker_input / use_chunk_plucker_post_attn.
  • cam_pos_embeds — Embeddings posicionais de câmera pré-computados opcionais (entradas "absmap" / "P") reutilizados em vez de recalculá-los.
  • pos_embeds — Incorporações de posição rotativa pré-computadas opcionais; quando None, elas são construídas a partir da forma do latente.
  • raymats — Matrizes de raios UCPE pré-computadas opcionais (usadas apenas quando cam_pos_embeds não carrega "P").
  • frame_valid_mask — Máscara de validade por quadro opcional, com forma (B, 1, T, 1, 1), (B, 1, T) ou (B, T).
  • chunk_size — Substituição do comprimento do chunk encaminhado aos blocos; caso não seja definido, utiliza o chunk_size configurado de cada bloco.

Execute o SANA-WM DiT.

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 )

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