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

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