Diffusers documentation
StableAudio3DiTModel
StableAudio3DiTModel
音声生成に使用される、Stable Audio 3における整流フロー速度予測型 Diffusion Transformer (DiT) です。
各 StableAudio3DiTBlock は以下を実行します:
- 自己注意機構 — 回転位置埋め込み(RoPE)付きの微分マルチヘッド注意機構。
- Cross-attention — T5Gemma テキストエンコーダーのトークンシーケンスにアテンションを適用します。
- フィードフォワード — SwiGLU 射影。
モデルは、timestep(指数関数Fourier特徴量 → 線形射影)と、StableAudio3DurationEmbedderからの長さ埋め込みであるグローバル条件ベクトルによって条件付けられます。
StableAudio3DiTModel
class diffusers.StableAudio3DiTModel
< source >( io_channels: int = 256patch_size: int = 1embed_dim: int = 1536depth: int = 24num_heads: int = 24cond_token_dim: int = 768global_cond_dim: int = 768local_add_cond_dim: int = 257timestep_features_dim: int = 256ff_mult: float = 4.0num_memory_tokens: int = 64use_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) — フーリエ時間ステップ特徴の出力次元数。 - 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
< source >( hidden_states: Tensortimestep: Tensorencoder_hidden_states: Tensorglobal_hidden_states: Tensorencoder_attention_mask: typing.Optional[torch.Tensor] = Nonelocal_add_cond: typing.Optional[torch.Tensor] = Nonereturn_dict: bool = True ) → StableAudio3DiTModelOutput or tuple
パラメータ
- 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
class diffusers.models.transformers.transformer_stable_audio3.StableAudio3DiTBlock
< source >( dim: intcontext_dim: intdim_heads: int = 64use_differential: bool = Trueff_mult: float = 4.0local_add_cond_dim: int = 257 )
単一のSA3 DiT transformerブロック。
処理順序:
- AdaLNモジュレーション付き自己注意(部分RoPE、RMS QK-norm)
- テキスト/長さコンテキストへのクロスアテンション(プレーンな RMS 事前正規化)
- 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
class diffusers.models.transformers.transformer_stable_audio3.StableAudio3DiTModelOutput
< source >( sample: Tensor )
StableAudio3DiTModelの出力です。