Hugging Face's logo

Diffusers documentation

StableAudio3DiTModel

You are viewing main version, which requires installation from source. If you'd like regular pip install, checkout the latest stable version (v0.39.0).
Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

StableAudio3DiTModel

A rectified-flow velocity-prediction Diffusion Transformer (DiT) for audio generation, used in Stable Audio 3.

Each StableAudio3DiTBlock performs:

  1. Self-attention — differential multi-head attention with rotary position embeddings (RoPE).
  2. Cross-attention — attends to the token sequence from the T5Gemma text encoder.
  3. 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

< >

( 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 of StableAudio3DiTBlock layers.
  • 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 to True) — 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 by to_cond_embed.
  • The global duration embedding (global_hidden_states) is projected by to_global_embed, summed with the timestep embedding, then expanded by global_cond_embedder into the per-block AdaLN modulation.
  • local_add_cond (inpainting) is projected per-block by to_local_embed.

num_memory_tokens learnable tokens are prepended to the audio sequence inside the transformer and removed before the output projection.

forward

< >

( 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 to True) — Whether to return a StableAudio3DiTModelOutput.

Returns

StableAudio3DiTModelOutput or tuple

the predicted velocity field, same shape as hidden_states.

StableAudio3DiTBlock

class diffusers.models.transformers.transformer_stable_audio3.StableAudio3DiTBlock

< >

( 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:

  1. AdaLN-modulated self-attention (partial RoPE, RMS QK-norm)
  2. Cross-attention to the text/duration context (plain RMS pre-norm)
  3. 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

< >

( sample: Tensor )

Parameters

  • sample (torch.Tensor) — The predicted velocity field, of the same shape as the input hidden_states.

The output of StableAudio3DiTModel.

Update on GitHub

←SkyReelsV2Transformer3DModel StableAudioDiTModel→