Diffusers documentation
StableAudio3DiTModel
StableAudio3DiTModel
Un Diffusion Transformer (DiT) de prédiction de vitesse à flux rectifié pour la génération audio, utilisé dans Stable Audio 3.
Chaque StableAudio3DiTBlock effectue :
- Auto-attention — attention multi-têtes différentielle avec des embeddings de position rotatifs (RoPE).
- Attention croisée — porte son attention sur la séquence de jetons issue de l'encodeur de texte T5Gemma.
- Feed-forward — projection SwiGLU.
Le modèle est conditionné sur un timestep (fonctions de Fourier exponentielles → projection linéaire) et un vecteur de conditionnement global (embedding de durée issu de StableAudio3DurationEmbedder).
StableAudio3DiTModel
classe 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 )
Paramètres
- io_channels (
int, par défaut 256) — Nombre de canaux latents. - patch_size (
int, par défaut 1) — Taille du patch temporel appliquée avant le transformer. - embed_dim (
int, par défaut 1536) — Dimension cachée du Transformer. - depth (
int, par défaut 24) — Nombre de couchesStableAudio3DiTBlock. - num_heads (
int, par défaut 24) — Nombre de têtes d'attention. - cond_token_dim (
int, par défaut 768) — Dimension des jetons de contexte de l'attention croisée. - global_cond_dim (
int, par défaut 768) — Dimension de l'embedding de durée globale. - local_add_cond_dim (
int, par défaut 257) — Nombre de canaux du tenseur additif local (inpaint). - timestep_features_dim (
int, par défaut 256) — Dimension de sortie des caractéristiques de pas de temps de Fourier. - ff_mult (
float, par défaut 4.0) — Facteur d'expansion du réseau de propagation avant SwiGLU. - num_memory_tokens (
int, par défaut 64) — Nombre de jetons mémoire apprenables. - use_differential_attention (
bool, par défautTrue) — Active l'auto-attention et l'attention croisée différentielles.
Le backbone Diffusion Transformer (DiT) de Stable Audio 3.
Le modèle prend en entrée un lot de latents audio bruités, un pas de temps scalaire, un contexte d'attention croisée (jetons de texte et de durée projetés) et un plongement de durée global, et prédit le champ de vitesse (objectif de flux rectifié).
Conditionnement :
- Le contexte d'attention croisée (
encoder_hidden_states) est projeté parto_cond_embed. - L'embedding de durée globale (
global_hidden_states) est projeté parto_global_embed, additionné à l'embedding de pas de temps, puis étendu parglobal_cond_embedderen modulation AdaLN par bloc. local_add_cond(inpainting) est projeté par bloc parto_local_embed.
num_memory_tokens jetons apprenables sont ajoutés au début de la séquence audio à l’intérieur du transformer et retirés avant la projection de sortie.
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
Paramètres
- hidden_states (
torch.Tensor) — Audio latent bruité(batch, io_channels, T). - timestep (
torch.Tensor) — Pas de diffusion(batch,)dans[0, 1]. - encoder_hidden_states (
torch.Tensor) — Contexte d'attention croisée(batch, T_ctx, cond_token_dim). - global_hidden_states (
torch.Tensor) — Embedding de durée globale(batch, global_cond_dim). - encoder_attention_mask (
torch.Tensor, facultatif) — Masque booléen(batch, T_ctx),True= valide. - local_add_cond (
torch.Tensor, facultatif) — Conditionnement additif local (inpaint)(batch, local_add_cond_dim, T). - return_dict (
bool, par défautTrue) — Indique s'il faut retourner unStableAudio3DiTModelOutput.
Retourne
StableAudio3DiTModelOutput ou tuple
le champ de vitesse prédit, de même forme que hidden_states.
StableAudio3DiTBlock
classe 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 )
Bloc de transformateur SA3 DiT unique.
Ordre des opérations :
- Auto-attention auto-régressive modulée par AdaLN (RoPE partiel, normalisation RMS QK)
- Attention croisée vers le contexte texte/durée (pré-norm RMS simple)
- Réseau de neurones à avant (feed-forward) SwiGLU modulé par AdaLN
La modulation AdaLN est to_scale_shift_gate + global_modulation, divisée en six blocs (scale_attn, shift_attn, gate_attn, scale_ff, shift_ff, gate_ff). Chaque branche à porte est mise à l'échelle par sigmoid(1 - gate). L'attention croisée n'est pas modulée par AdaLN, conformément à la référence (cross_attend_norm est une simple norme RMS).
Lorsque local_seq est fourni (inpainting), il est projeté par bloc par to_local_embed et ajouté aux positions audio du flux résiduel après l'attention croisée (et avant le réseau de neurones feed-forward), conformément à la référence.
StableAudio3DiTModelOutput
classe diffusers.models.transformers.transformer_stable_audio3.StableAudio3DiTModelOutput
< source >( sample: Tensor )
La sortie de StableAudio3DiTModel.