Buckets:

hf-doc-build/doc / diffusers /main /ja /api /models /stable_audio_3_transformer.md
mishig's picture
|
download
raw
7.95 kB

StableAudio3DiTModel

音声生成に使用される、Stable Audio 3における整流フロー速度予測型 Diffusion Transformer (DiT) です。

各 StableAudio3DiTBlock は以下を実行します:

  1. 自己注意機構 — 回転位置埋め込み(RoPE)付きの微分マルチヘッド注意機構。
  2. Cross-attention — T5Gemma テキストエンコーダーのトークンシーケンスにアテンションを適用します。
  3. フィードフォワード — SwiGLU 射影。

モデルは、timestep(指数関数Fourier特徴量 → 線形射影)と、StableAudio3DurationEmbedderからの長さ埋め込みであるグローバル条件ベクトルによって条件付けられます。

StableAudio3DiTModel[[diffusers.StableAudio3DiTModel]]

diffusers.StableAudio3DiTModel[[diffusers.StableAudio3DiTModel]]

diffusers.StableAudio3DiTModel(io_channels: int = 256, patch_size: int = 1, embed_dim: int = 1536, depth: int = 24, num_heads: int = 24, cond_token_dim: int = 768, global_cond_dim: int = 768, local_add_cond_dim: int = 257, timestep_features_dim: int = 256, ff_mult: float = 4.0, num_memory_tokens: int = 64, use_differential_attention: bool = True)

ソース

パラメータ:

io_channels (int、デフォルトは 256) : 潜在チャネルの数。

patch_size (int、デフォルトは 1) : Transformer に適用される時間方向のパッチサイズ。

embed_dim (int、デフォルトは 1536) : Transformer の隠れ層の次元数。

depth (int、デフォルトは 24) : StableAudio3DiTBlock レイヤーの数。

num_heads (int、デフォルトは 24) : アテンションヘッドの数。

cond_token_dim (int、デフォルトは 768) : クロスアテンションのコンテキストトークンの次元数。

global_cond_dim (int、デフォルトは 768) : グローバルな期間埋め込みの次元数。

local_add_cond_dim (int、デフォルトは 257) : ローカル加算型(インペイント)テンソルのチャネル数。

timestep_features_dim (int、デフォルトは 256) : Fourier 時間ステップ特徴の出力次元。

ff_mult (float、デフォルトは 4.0) : SwiGLU フィードフォワードの拡張係数。

num_memory_tokens (int、デフォルトは 64) : 学習可能なメモリトークンの数。

use_differential_attention (bool、デフォルトはTrue) : 差分自己/クロスアテンションを有効にします。

Stable Audio 3のDiffusion Transformer (DiT)バックボーン。

モデルは、ノイズ付きオーディオ潜在変数のバッチ、スカラーのタイムステップ、クロスアテンションコンテキスト(射影されたテキストおよび期間トークン)、グローバルな期間埋め込みを受け取り、速度場を予測します(rectified-flow目的関数)。

条件付け:

  • Cross-attention のコンテキスト(encoder_hidden_states)は、to_cond_embed によって射影されます。
  • グローバル期間埋め込み(global_hidden_states)はto_global_embedによって射影され、タイムステップ埋め込みと加算された後、global_cond_embedderによってブロックごとのAdaLNモジュレーションに展開されます。
  • インペインティング用のlocal_add_condは、to_local_embedによってブロックごとに射影されます。

num_memory_tokens の学習可能トークンは、transformer 内で音声シーケンスの先頭に付加され、出力投影の前に削除されます。

forward[[diffusers.StableAudio3DiTModel.forward]]

forward(hidden_states: Tensor, timestep: Tensor, encoder_hidden_states: Tensor, global_hidden_states: Tensor, encoder_attention_mask: typing.Optional[torch.Tensor] = None, local_add_cond: typing.Optional[torch.Tensor] = None, return_dict: bool = True)

ソース

パラメータ:

hidden_states (torch.Tensor) : ノイズ付き潜在オーディオ (batch, io_channels, T)。

timestep (torch.Tensor) : [0, 1] 範囲の Diffusion timestep (batch,)。

encoder_hidden_states (torch.Tensor) : クロスアテンションのコンテキスト (batch, T_ctx, cond_token_dim)。

global_hidden_states (torch.Tensor) : グローバル期間埋め込み (batch, global_cond_dim)。

encoder_attention_mask (torch.Tensor, 任意) : Bool 型のマスク (batch, T_ctx)。True は有効を意味します。

local_add_cond (torch.Tensor, 任意) : ローカル加算(インペイント)条件付け (batch, local_add_cond_dim, T)。

return_dict (bool、デフォルトはTrue) : StableAudio3DiTModelOutputを返すかどうか。

戻り値: StableAudio3DiTModelOutput または tuple

hidden_states と同じ形状の予測速度場。

StableAudio3DiTBlock[[diffusers.models.transformers.transformer_stable_audio3.StableAudio3DiTBlock]]

diffusers.models.transformers.transformer_stable_audio3.StableAudio3DiTBlock[[diffusers.models.transformers.transformer_stable_audio3.StableAudio3DiTBlock]]

diffusers.models.transformers.transformer_stable_audio3.StableAudio3DiTBlock(dim: int, context_dim: int, dim_heads: int = 64, use_differential: bool = True, ff_mult: float = 4.0, local_add_cond_dim: int = 257)

ソース

単一のSA3 DiT transformerブロック。

処理順序:

  1. AdaLNモジュレーション付き自己注意(部分RoPE、RMS QK-norm)
  2. テキスト/長さコンテキストへのクロスアテンション(プレーンな RMS 事前正規化)
  3. AdaLNモジュレーション付きSwiGLUフィードフォワード

AdaLN 変調は to_scale_shift_gate + global_modulation であり、(scale_attn, shift_attn, gate_attn, scale_ff, shift_ff, gate_ff) の 6 つのチャンクに分割されます。各ゲーティングされたブランチは sigmoid(1 - gate) によってスケーリングされます。クロスアテンションは AdaLN 変調が適用されていません。これはリファレンス実装と一致しており、cross_attend_norm は通常の RMS 正規化です。

local_seq が指定されている場合(inpainting)、to_local_embed によってブロックごとに射影され、クロスアテンションの後にフィードフォワードの前に、残差ストリームのオーディオ位置に追加されます。これは参照実装と一致します。

StableAudio3DiTModelOutput[[diffusers.models.transformers.transformer_stable_audio3.StableAudio3DiTModelOutput]]

diffusers.models.transformers.transformer_stable_audio3.StableAudio3DiTModelOutput[[diffusers.models.transformers.transformer_stable_audio3.StableAudio3DiTModelOutput]]

diffusers.models.transformers.transformer_stable_audio3.StableAudio3DiTModelOutput(sample: Tensor)

ソース

パラメータ:

sample (torch.Tensor) : 入力 hidden_states と同じ形状の、予測された速度場です。

StableAudio3DiTModelの出力です。

Xet Storage Details

Size:
7.95 kB
·
Xet hash:
8a8c48f35e7950f7b4e049eaa4915bfba29e782a7cecdf24e0dafd7b73589dff

Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.