Hugging Face's logo

Diffusers documentation

SanaWMTransformer3DModel

Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

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

< >

( 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

< >

( 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.

unpatchify

< >

( xpost_patch_num_framespost_patch_heightpost_patch_width )

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

Transformer2DModelOutput

класс diffusers.models.modeling_outputs.Transformer2DModelOutput

< >

( 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.

Update on GitHub

←SanaWMLTX2RefinerTransformer3DModel SD3Transformer2DModel→