Diffusers documentation
SanaWMTransformer3DModel
SanaWMTransformer3DModel
Un Diffusion Transformer 3D (1.6B parámetros) para la generación de video a partir de imágenes con control de cámara, utilizado como muestreador de la etapa 1 de SanaWMPipeline. El transformer combina:
- una rama principal de atención lineal GDN-Triton bidireccional (profundidad 20, oculto 2240, 20 cabezas),
- una rama de control de cámara UCPE (Unified Camera Pose Embedding) que consume una representación raymap + Plücker de la trayectoria solicitada, y
- un incrustado de posición rotacional 3D de estilo Wan más bloques de atención softmax periódicos inyectados cada
softmax_every_ncapas.
El modelo se puede cargar con:
import torch
from diffusers import SanaWMTransformer3DModel
transformer = SanaWMTransformer3DModel.from_pretrained(
"Efficient-Large-Model/SANA-WM_bidirectional-diffusers",
subfolder="transformer",
torch_dtype=torch.bfloat16,
)SanaWMTransformer3DModel
clase 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, por defecto 128) — Canales latentes del VAE (LTX-2).
- softmax_every_n (int, valor predeterminado 4) — Utiliza un bloque de atención softmax cada N bloques.
- linear_head_dim (int, por defecto 112) — Dimensión de la cabeza GDN.
- ffn_type (str, por defecto
"GLUMBConvTemp") — FFN. - t_kernel_size (int, por defecto 3) — Núcleo de convolución temporal.
- conv_kernel_size (int, valor predeterminado 4) — Núcleo de convolución espacial dentro de la atención.
- k_conv_only (bool, por defecto True) — Aplicar conv solo en K.
- pos_embed_type (str, por defecto
"wan_rope") — Incrustación de posición. - qk_norm (bool, por defecto True) — RMSNorm en Q/K.
- cross_norm (bool, por defecto True) — RMSNorm en K de la atención cruzada.
- y_norm (bool, por defecto True) — Aplica
attention_y_norma las incrustaciones de texto. - init_cam_from_base (bool, predeterminado en True) — Sin uso; la rama de la cámara se carga desde el punto de control. Se mantiene para que los archivos config.json publicados se carguen.
- use_chunk_plucker_post_attn (bool, por defecto True). —
- chunk_plucker_channels (int, valor predeterminado 48) —
6 dims * temporal_stride 8. - chunk_plucker_post_attn_blocks (int, por defecto 20) — Todos los bloques.
- fp32_attention (bool, predeterminado en True) — Sin uso; la atención siempre se ejecuta en fp32. Se mantiene para que los archivos config.json publicados se carguen.
- image_size (int, por defecto 720) — Tamaño nominal de la imagen.
- caption_channels (int, valor predeterminado 2304) — Tamaño oculto de Gemma-2.
SANA-WM 1600M DiT bidireccional controlado por cámara.
Un DiT de una sola clase (depth=20, hidden_size=2240, patch_size=(1,1,1), num_heads=20, es decir, la publicación pública Efficient-Large-Model/SANA-WM_bidirectional). save_pretrained / from_pretrained funcionan directamente mediante ConfigMixin.
Cada bloque ejecuta una atención condicionada por cámara (UCPE): [BidirectionalGDNUCPESinglePathLiteLA], excepto cada bloque softmax_every_n-ésimo, que ejecuta su contraparte de softmax [_SoftmaxUCPESinglePathLiteLA].
El state-dict es idéntico al checkpoint público de sana, excepto por el buffer pos_embed que se eliminó 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)pasos de difusión por cuadro (estilo LTX). - encoder_hidden_states —
(B, 1, L, D_caption)incrustaciones de texto. - encoder_attention_mask — máscara de atención de texto
(B, L)(convención de diffusers). - mask — Alias de
encoder_attention_maskque coincide con el nombre del kwarg del DiT de sana. Si se pasan ambos,masktiene prioridad. - return_dict — Si es
True(por defecto), devuelve unTransformer2DModelOutput; en caso contrario, devuelve una tupla de un elemento(sample,). - data_info — Condicionamiento adicional;
data_info["image_vae_embeds"]se concatena con los latentes a lo largo del eje de canales cuando está presente. - camera_conditions —
(B, T, 20)condiciones de cámara en bruto que controlan la rama de control de cámara (UCPE). - chunk_plucker — Incrustaciones de rayos Plucker
(B, C, T, H, W), consumidas cuando el modelo está configurado conuse_chunk_plucker_input/use_chunk_plucker_post_attn. - cam_pos_embeds — Embeddings posicionales de cámara precalculados opcionales
(entradas de
"absmap"/"P") que se reutilizan en lugar de recalcularlos. - pos_embeds — Incrustaciones de posición rotatoria precalculadas opcionales; cuando es
None, se construyen a partir de la forma latente. - raymats — Matrices de rayos UCPE precalculadas opcionales (se usan solo cuando
cam_pos_embedsno lleva"P"). - frame_valid_mask — Máscara de validez por cuadro opcional, con forma
(B, 1, T, 1, 1),(B, 1, T)o(B, T). - chunk_size — Anulación de la longitud del fragmento que se reenvía a los bloques; si no se especifica, se utiliza el
chunk_sizeconfigurado en cada bloque.
Ejecuta el DiT de SANA-WM.
x: (N, T, patch_size*2 C) imgs: (N, H, W, C)
Transformer2DModelOutput
clase diffusers.models.modeling_outputs.Transformer2DModelOutput
< source >( sample: torch.Tensor )
Parámetros
- sample (
torch.Tensorde 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 entradaencoder_hidden_states. Si es discreto, devuelve distribuciones de probabilidad para los píxeles latentes sin ruido.
La salida de Transformer2DModel.