Diffusers documentation
SanaWMTransformer3DModel
SanaWMTransformer3DModel
3D-диффузионный трансформер (1,6 млрд параметров) для генерации видео из изображения с управлением камерой, используемый в качестве сэмплера первого этапа в SanaWMPipeline. Трансформер объединяет:
- двунаправленная основная ветвь линейного внимания GDN-Triton (глубина 20, скрытое пространство 2240, 20 головок),
- ветвь управления камерой UCPE (Unified Camera Pose Embedding), которая принимает raymap + представление Plücker для запрошенной траектории, и
- 3D-вращательное позиционное встраивание в стиле Wan плюс блоки периодического softmax-внимания, вставляемые каждые
softmax_every_nслои.
Модель можно загрузить с помощью:
import torch
from diffusers import SanaWMTransformer3DModel
transformer = SanaWMTransformer3DModel.from_pretrained(
"Efficient-Large-Model/SANA-WM_bidirectional-diffusers",
subfolder="transformer",
torch_dtype=torch.bfloat16,
)SanaWMTransformer3DModel
класс 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' )
Параметры
- in_channels (int, по умолчанию 128) — каналы латентного пространства VAE (LTX-2).
- softmax_every_n (int, по умолчанию 4) — Использовать блок внимания с softmax каждые N блоков.
- linear_head_dim (int, по умолчанию 112) — размерность головы GDN.
- ffn_type (str, по умолчанию
"GLUMBConvTemp") — FFN. - t_kernel_size (int, по умолчанию 3) — Ядро временной свёртки.
- conv_kernel_size (int, по умолчанию 4) — размер пространственного свёрточного ядра внутри механизма внимания.
- k_conv_only (bool, по умолчанию True) — Применять свёртку только к K.
- pos_embed_type (str, по умолчанию
"wan_rope") — Встраивание позиции. - qk_norm (bool, по умолчанию True) — RMSNorm для Q/K.
- cross_norm (bool, по умолчанию True) — RMSNorm для K в кросс-внимании.
- y_norm (bool, по умолчанию True) — применяет
attention_y_normк текстовым эмбеддингам. - init_cam_from_base (bool, по умолчанию True) — Не используется; ветка камеры загружается из контрольной точки. Оставлено для корректной загрузки опубликованных файлов config.json.
- use_chunk_plucker_post_attn (bool, по умолчанию True). —
- chunk_plucker_channels (int, по умолчанию 48) —
6 dims * temporal_stride 8. - chunk_plucker_post_attn_blocks (int, по умолчанию 20) — Все блоки.
- fp32_attention (bool, по умолчанию True) — Не используется; внимание всегда выполняется в fp32. Оставлено для корректной загрузки опубликованных файлов config.json.
- image_size (int, по умолчанию 720) — Номинальный размер изображения.
- caption_channels (int, по умолчанию 2304) — размер скрытого слоя Gemma-2.
SANA-WM 1600M — двунаправленный DiT с управлением камерой.
Одноклассная DiT (depth=20, hidden_size=2240, patch_size=(1,1,1), num_heads=20 — то есть публичный релиз Efficient-Large-Model/SANA-WM_bidirectional). save_pretrained / from_pretrained работают из коробки благодаря ConfigMixin.
Каждый блок выполняет внимание с условием камеры (UCPE): [BidirectionalGDNUCPESinglePathLiteLA], за исключением каждого softmax_every_n-го блока, который выполняет его softmax-аналог [_SoftmaxUCPESinglePathLiteLA].
Файл state-dict идентичен публичному чекпоинту sana, за исключением намеренно удалённого буфера pos_embed.
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 )
Параметры
- hidden_states — латенты
(B, C, T, H, W). - timestep —
(B, 1, T)временные шаги диффузии для каждого кадра (стиль LTX). - encoder_hidden_states — текстовые эмбеддинги
(B, 1, L, D_caption). - encoder_attention_mask — маска внимания для текста
(B, L)(конвенция diffusers). - mask — Псевдоним для
encoder_attention_mask, соответствующий имени параметра kwarg в sana DiT. Если переданы оба,maskимеет приоритет. - return_dict — если
True(по умолчанию), возвращаетTransformer2DModelOutput; иначе возвращает кортеж из одного элемента(sample,). - data_info — дополнительные условия;
data_info["image_vae_embeds"]конкатенируется с латентами по оси каналов, если присутствует. - camera_conditions —
(B, T, 20)исходные условия камеры, управляющие ветвью управления камерой (UCPE). - chunk_plucker — Вложения лучей Plucker
(B, C, T, H, W), используются, когда модель настроена сuse_chunk_plucker_input/use_chunk_plucker_post_attn. - cam_pos_embeds — Необязательные предварительно вычисленные позиционные эмбеддинги камеры (записи
"absmap"/"P"), которые используются повторно вместо их пересчёта. - pos_embeds — необязательные предварительно вычисленные вращательные позиционные эмбеддинги; если
None, они создаются на основе формы латентов. - raymats — необязательные предварительно вычисленные матрицы лучей UCPE (используются только в том случае, если
cam_pos_embedsне содержит"P"). - frame_valid_mask — Необязательная маска валидности для каждого кадра, с формой
(B, 1, T, 1, 1),(B, 1, T)или(B, T). - chunk_size — Переопределение длины чанка, передаваемое блокам; при отсутствии используется настроенное для каждого блока значение
chunk_size.
Запустите SANA-WM DiT.
x: (N, T, patch_size*2 C) imgs: (N, H, W, C)
Transformer2DModelOutput
класс diffusers.models.modeling_outputs.Transformer2DModelOutput
< source >( sample: torch.Tensor )
Параметры
- sample (
torch.Tensorформы(batch_size, num_channels, height, width)или(batch size, num_vector_embeds - 1, num_latent_pixels), если Transformer2DModel дискретный) — Выходные скрытые состояния, обусловленные входомencoder_hidden_states. Если модель дискретная, возвращает распределения вероятностей для латентных пикселей без шума.
Вывод Transformer2DModel.