Diffusers documentation
SanaWMTransformer3DModel
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_ncamadas.
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
< source >( 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_normaos 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
< source >( 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_maskcorrespondente ao nome do kwarg do DiT do sana. Se ambos forem passados,masktem precedência. - return_dict — Se
True(padrão), retorna umTransformer2DModelOutput; 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 comuse_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_embedsnã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_sizeconfigurado de cada bloco.
Execute o SANA-WM DiT.
x: (N, T, patch_size*2 C) imgs: (N, H, W, C)
Transformer2DModelOutput
classe diffusers.models.modeling_outputs.Transformer2DModelOutput
< source >( sample: torch.Tensor )
Parâmetros
- sample (
torch.Tensorde 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 à entradaencoder_hidden_states. Se discreto, retorna distribuições de probabilidade para os pixels latentes sem ruído.
A saída de Transformer2DModel.