Diffusers documentation
SanaWMTransformer3DModel
SanaWMTransformer3DModel
カメラ制御による画像から動画への生成に使用される3次元拡散Transformer（16億パラメータ）であり、SanaWMPipelineのステージ1サンプラーとして機能します。このTransformerは以下を組み合わせます。
- 双方向 GDN-Triton 線形アテンションのメインブランチ（深さ 20、隠れ層サイズ 2240、ヘッド数 20）、
- UCPE（Unified Camera Pose Embedding）のカメラ制御ブランチ。これは、要求された軌道のレイマップとPlücker表現を消費し、
- Wanスタイルの3D回転位置埋め込みと、
softmax_every_n層ごとに挿入される周期的なsoftmaxアテンションブロック。
モデルは以下のように読み込むことができます:
import torch
from diffusers import SanaWMTransformer3DModel
transformer = SanaWMTransformer3DModel.from_pretrained(
"Efficient-Large-Model/SANA-WM_bidirectional-diffusers",
subfolder="transformer",
torch_dtype=torch.bfloat16,
)SanaWMTransformer3DModel
class 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) — N ブロックごとに 1 つの softmax attention ブロックを使用します。
- linear_head_dim (int、デフォルト値は 112) — GDN のヘッド次元。
- ffn_type (str、デフォルトは
"GLUMBConvTemp") — FFN。 - t_kernel_size (int、デフォルトは 3) — 時間方向の畳み込みカーネルサイズ。
- conv_kernel_size (int、デフォルトは 4) — attention 内の空間的な conv カーネルサイズ。
- k_conv_only (bool、デフォルトは True) — conv を K にのみ適用します。
- pos_embed_type (str、デフォルトは
"wan_rope") — 位置埋め込み。 - qk_norm (bool、デフォルトは True) — Q/K に対する RMSNorm。
- cross_norm (bool、デフォルトは True) — クロスアテンションの K に対する RMSNorm。
- 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) — 未使用。attention は常に 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 は、意図的に削除された pos_embed バッファを除き、公開されている sana チェックポイントと同一です。
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の別名です。sana DiT の kwarg 名に一致します。両方が渡された場合、maskが優先されます。 - return_dict —
True（デフォルト）の場合、Transformer2DModelOutputを返します。それ以外の場合は、1要素のタプル(sample,)を返します。 - data_info — 追加の条件付け情報。
data_info["image_vae_embeds"]が存在する場合、チャネル軸に沿って潜在変数に連結されます。 - camera_conditions — カメラ制御（UCPE）ブランチを駆動する生のカメラス状態。
(B, T, 20) - 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)、または Transformer2DModel が離散型の場合(batch size, num_vector_embeds - 1, num_latent_pixels)) —encoder_hidden_states入力を条件とする隠れ状態の出力です。離散型の場合、ノイズ除去されていない潜在ピクセルの確率分布を返します。
Transformer2DModelの出力です。