Diffusers documentation
StableAudio3DiTModel
StableAudio3DiTModel
Um Diffusion Transformer (DiT) de previsão de velocidade de fluxo retificado para geração de áudio, usado no Stable Audio 3.
Cada StableAudio3DiTBlock executa:
- Autoatenção — atenção multi-cabeça diferencial com incorporações de posição rotativas (RoPE).
- Atenção cruzada — presta atenção à sequência de tokens do codificador de texto T5Gemma.
- Feed-forward — projeção SwiGLU.
O modelo é condicionado em um timestep (recursos de Fourier exponenciais → projeção linear) e em um vetor de condicionamento global (embedding de duração do 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 )
Parâmetros
- io_channels (
int, padrão 256) — Número de canais latentes. - patch_size (
int, padrão 1) — Tamanho do patch temporal aplicado antes do transformer. - embed_dim (
int, padrão 1536) — Dimensão oculta do Transformer. - depth (
int, padrão 24) — Número de camadasStableAudio3DiTBlock. - num_heads (
int, padrão 24) — Número de cabeças de atenção. - cond_token_dim (
int, padrão 768) — Dimensão dos tokens de contexto da atenção cruzada. - global_cond_dim (
int, padrão 768) — Dimensão do embedding de duração global. - local_add_cond_dim (
int, padrão 257) — Canais do tensor aditivo local (inpaint). - timestep_features_dim (
int, padrão 256) — Dimensão de saída dos recursos de timestep Fourier. - ff_mult (
float, padrão 4.0) — Fator de expansão do feed-forward SwiGLU. - num_memory_tokens (
int, padrão 64) — Número de tokens de memória aprendíveis. - use_differential_attention (
bool, padrãoTrue) — Ativa a auto/atenção diferencial.
O backbone Diffusion Transformer (DiT) do Stable Audio 3.
O modelo recebe um lote de latentes de áudio com ruído, um timestep escalar, um contexto de cross-attention (tokens de texto e duração projetados) e um embedding global de duração, e prevê o campo de velocidade (objetivo de fluxo retificado).
Condição:
- O contexto de atenção cruzada (
encoder_hidden_states) é projetado porto_cond_embed. - O embedding de duração global (
global_hidden_states) é projetado porto_global_embed, somado ao embedding de timestep e, em seguida, expandido porglobal_cond_embedderpara a modulação AdaLN por bloco. local_add_cond(inpainting) é projetado por bloco porto_local_embed.
num_memory_tokens tokens aprendíveis são adicionados ao início da sequência de áudio dentro do transformer e removidos antes da projeção de saída.
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
Parâmetros
- hidden_states (
torch.Tensor) — Áudio latente com ruído(batch, io_channels, T). - timestep (
torch.Tensor) — Timestep de difusão(batch,)em[0, 1]. - encoder_hidden_states (
torch.Tensor) — Contexto de atenção cruzada(batch, T_ctx, cond_token_dim). - global_hidden_states (
torch.Tensor) — Embedding de duração global(batch, global_cond_dim). - encoder_attention_mask (
torch.Tensor, opcional) — Máscara Bool(batch, T_ctx),True= válido. - local_add_cond (
torch.Tensor, opcional) — Condição aditiva local (inpaint)(batch, local_add_cond_dim, T). - return_dict (
bool, padrãoTrue) — Se deve retornar umStableAudio3DiTModelOutput.
Retorna
StableAudio3DiTModelOutput ou tuple
o campo de velocidade previsto, com a mesma forma de 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 )
Bloco único de transformador DiT SA3.
Ordem de operações:
- Autoatenção modulada por AdaLN (RoPE parcial, normalização RMS QK)
- Atenção cruzada para o contexto de texto/duração (pré-norm RMS simples)
- Feed-forward SwiGLU modulado por AdaLN
A modulação AdaLN é to_scale_shift_gate + global_modulation, dividida em seis blocos (scale_attn, shift_attn, gate_attn, scale_ff, shift_ff, gate_ff). Cada ramo com gate é escalado por sigmoid(1 - gate). A atenção cruzada não é modulada por AdaLN, em conformidade com a referência (cross_attend_norm é uma normalização RMS simples).
Quando local_seq é fornecido (inpainting), ele é projetado por bloco por to_local_embed e adicionado às posições de áudio do fluxo residual após a atenção cruzada (e antes do feed-forward), em conformidade com a referência.
StableAudio3DiTModelOutput
classe diffusers.models.transformers.transformer_stable_audio3.StableAudio3DiTModelOutput
< source >( sample: Tensor )
A saída de StableAudio3DiTModel.