Diffusers documentation
StableAudio3DiTModel
StableAudio3DiTModel
用于音频生成的整流流速度预测扩散 Transformer（DiT），应用于 Stable Audio 3。
每个 StableAudio3DiTBlock 执行以下操作：
- 自注意力 — 带有旋转位置编码（RoPE）的微分多头注意力。
- 交叉注意力 — 关注来自 T5Gemma 文本编码器的 token 序列。
- 前馈 — SwiGLU 投影。
该模型以 timestep（指数傅里叶特征 → 线性投影）和 全局条件向量（来自 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 之前应用的时间维度 patch 大小。 - 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) — 局部加性（inpaint）张量的通道数。 - timestep_features_dim (
int，默认为 256) — 傅里叶时间步特征的输出维度。 - ff_mult (
float，默认值为 4.0) — SwiGLU 前馈扩展因子。 - num_memory_tokens (
int，默认为 64) — 可学习记忆 token 的数量。 - use_differential_attention（
bool，默认为True）— 启用差分自注意力/交叉注意力。
Stable Audio 3 的 Diffusion Transformer (DiT) 骨干网络。
模型接收一批带噪声的音频潜变量、一个标量时间步、交叉注意力上下文（投影后的文本和时长 token）以及一个全局时长嵌入，并预测速度场（rectified-flow 目标）。
条件设置：
- 交叉注意力上下文（
encoder_hidden_states）由to_cond_embed进行投影。 - 全局时长嵌入（
global_hidden_states）由to_global_embed进行投影，与时间步嵌入相加，然后通过global_cond_embedder扩展为每个块的 AdaLN 调制。 local_add_cond（inpainting）通过to_local_embed按块进行投影。
num_memory_tokens 个可学习 token 会在 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) — 扩散时间步(batch,)，取值范围为[0, 1]。 - 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，可选) — 布尔掩码(batch, T_ctx)，True表示有效。 - local_add_cond (
torch.Tensor, 可选) — 局部加性（inpaint）条件(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 变换器块。
操作顺序：
- AdaLN 调制自注意力（部分 RoPE，RMS QK 归一化）
- 针对文本/时长上下文的交叉注意力（普通 RMS 前置归一化）
- AdaLN 调制的 SwiGLU 前馈
AdaLN 调制为 to_scale_shift_gate + global_modulation，拆分为六个块 (scale_attn, shift_attn, gate_attn, scale_ff, shift_ff, gate_ff)。每个门控分支都会乘以 sigmoid(1 - gate) 进行缩放。交叉注意力不进行 AdaLN 调制，这与参考实现一致（cross_attend_norm 是普通的 RMS 归一化）。
当提供 local_seq（用于图像修复）时，它会通过 to_local_embed 按块进行投影，并添加到交叉注意力之后（前馈层之前）残差流中的音频位置，以与参考实现保持一致。
StableAudio3DiTModelOutput
class diffusers.models.transformers.transformer_stable_audio3.StableAudio3DiTModelOutput
< source >( sample: Tensor )
StableAudio3DiTModel 的输出。