Diffusers documentation
StableAudio3DiTModel
StableAudio3DiTModel
A rectified-flow velocity-prediction Diffusion Transformer (DiT) for audio generation, used in Stable Audio 3.
Each StableAudio3DiTBlock performs:
- Self-attention — differential multi-head attention with rotary position embeddings (RoPE).
- Cross-attention — attends to the token sequence from the T5Gemma text encoder.
- Feed-forward — SwiGLU projection.
The model is conditioned on a timestep (exponential Fourier features → linear projection) and a global conditioning vector (duration embedding from 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 )
Parameters
- io_channels (
int, defaults to 256) — Number of latent channels. - patch_size (
int, defaults to 1) — Temporal patch size applied before the transformer. - embed_dim (
int, defaults to 1536) — Transformer hidden dimension. - depth (
int, defaults to 24) — Number ofStableAudio3DiTBlocklayers. - num_heads (
int, defaults to 24) — Number of attention heads. - cond_token_dim (
int, defaults to 768) — Dimension of the cross-attention context tokens. - global_cond_dim (
int, defaults to 768) — Dimension of the global duration embedding. - local_add_cond_dim (
int, defaults to 257) — Channels of the local-additive (inpaint) tensor. - timestep_features_dim (
int, defaults to 256) — Output dimension of the Fourier timestep features. - ff_mult (
float, defaults to 4.0) — SwiGLU feed-forward expansion factor. - num_memory_tokens (
int, defaults to 64) — Number of learnable memory tokens. - use_differential_attention (
bool, defaults toTrue) — Enable differential self/cross attention.
The Diffusion Transformer (DiT) backbone of Stable Audio 3.
The model takes a batch of noisy audio latents, a scalar timestep, a cross-attention context (projected text and duration tokens), and a global duration embedding, and predicts the velocity field (rectified-flow objective).
Conditioning:
- Cross-attention context (
encoder_hidden_states) is projected byto_cond_embed. - The global duration embedding (
global_hidden_states) is projected byto_global_embed, summed with the timestep embedding, then expanded byglobal_cond_embedderinto the per-block AdaLN modulation. local_add_cond(inpainting) is projected per-block byto_local_embed.
num_memory_tokens learnable tokens are prepended to the audio sequence inside the transformer and removed before
the output projection.
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
Parameters
- hidden_states (
torch.Tensor) — Noisy latent audio(batch, io_channels, T). - timestep (
torch.Tensor) — Diffusion timestep(batch,)in[0, 1]. - encoder_hidden_states (
torch.Tensor) — Cross-attention context(batch, T_ctx, cond_token_dim). - global_hidden_states (
torch.Tensor) — Global duration embedding(batch, global_cond_dim). - encoder_attention_mask (
torch.Tensor, optional) — Bool mask(batch, T_ctx),True= valid. - local_add_cond (
torch.Tensor, optional) — Local-additive (inpaint) conditioning(batch, local_add_cond_dim, T). - return_dict (
bool, defaults toTrue) — Whether to return aStableAudio3DiTModelOutput.
Returns
StableAudio3DiTModelOutput or tuple
the predicted velocity field, same shape as 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 )
Single SA3 DiT transformer block.
Order of operations:
- AdaLN-modulated self-attention (partial RoPE, RMS QK-norm)
- Cross-attention to the text/duration context (plain RMS pre-norm)
- AdaLN-modulated SwiGLU feed-forward
The AdaLN modulation is to_scale_shift_gate + global_modulation, split into six chunks (scale_attn, shift_attn, gate_attn, scale_ff, shift_ff, gate_ff). Each gated branch is scaled by sigmoid(1 - gate). Cross-attention is not AdaLN-modulated, matching the reference (cross_attend_norm is a plain RMS norm).
When local_seq is provided (inpainting), it is projected per-block by to_local_embed and added to the audio
positions of the residual stream after cross-attention (and before the feed-forward), matching the reference.
StableAudio3DiTModelOutput
class diffusers.models.transformers.transformer_stable_audio3.StableAudio3DiTModelOutput
< source >( sample: Tensor )
The output of StableAudio3DiTModel.