Diffusers documentation
StableAudio3DiTModel
StableAudio3DiTModel
Un Diffusion Transformer (DiT) de predicción de velocidad de flujo rectificado para la generación de audio, utilizado en Stable Audio 3.
Cada StableAudio3DiTBlock realiza:
- Autoatención — atención multi-cabeza diferencial con incrustaciones de posición rotacional (RoPE).
- Atención cruzada — atiende a la secuencia de tokens del codificador de texto T5Gemma.
- Feed-forward — Proyección SwiGLU.
El modelo se condiciona en un timestep (características de Fourier exponenciales → proyección lineal) y un vector de condicionamiento global (incrustación de duración desde StableAudio3DurationEmbedder).
StableAudio3DiTModel
clase 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, por defecto 256) — Número de canales latentes. - patch_size (
int, por defecto 1) — Tamaño de parche temporal aplicado antes del transformer. - embed_dim (
int, por defecto 1536) — Dimensión oculta del Transformer. - depth (
int, por defecto 24) — Número de capasStableAudio3DiTBlock. - num_heads (
int, por defecto 24) — Número de cabezas de atención. - cond_token_dim (
int, por defecto 768) — Dimensión de los tokens de contexto de la atención cruzada. - global_cond_dim (
int, por defecto 768) — Dimensión del embedding de duración global. - local_add_cond_dim (
int, por defecto 257) — Canales del tensor local-aditivo (inpaint). - timestep_features_dim (
int, por defecto 256) — Dimensión de salida de las características de paso de tiempo de Fourier. - ff_mult (
float, valor predeterminado 4.0) — Factor de expansión del feed-forward de SwiGLU. - num_memory_tokens (
int, por defecto 64) — Número de tokens de memoria aprendibles. - use_differential_attention (
bool, por defectoTrue) — Habilita la autoatención/cruce diferencial.
El backbone Diffusion Transformer (DiT) de Stable Audio 3.
El modelo toma un lote de latentes de audio con ruido, un paso temporal escalar, un contexto de atención cruzada (tokens de texto y duración proyectados) y una incrustación de duración global, y predice el campo de velocidad (objetivo de flujo rectificado).
Condicionamiento:
- El contexto de atención cruzada (
encoder_hidden_states) se proyecta medianteto_cond_embed. - El embedding de duración global (
global_hidden_states) se proyecta medianteto_global_embed, se suma con el embedding de paso de tiempo y luego se expande medianteglobal_cond_embedderen la modulación AdaLN por bloque. local_add_cond(inpainting) se proyecta por bloque medianteto_local_embed.
Se añaden num_memory_tokens tokens aprendibles al inicio de la secuencia de audio dentro del transformer y se eliminan antes de la proyección de salida.
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) — Audio latente con ruido(batch, io_channels, T). - timestep (
torch.Tensor) — Paso de difusión(batch,)en[0, 1]. - encoder_hidden_states (
torch.Tensor) — Contexto de atención cruzada(batch, T_ctx, cond_token_dim). - global_hidden_states (
torch.Tensor) — Incrustación de duración global(batch, global_cond_dim). - encoder_attention_mask (
torch.Tensor, opcional) — Máscara booleana(batch, T_ctx),True= válido. - local_add_cond (
torch.Tensor, opcional) — Condicionamiento aditivo local (inpaint)(batch, local_add_cond_dim, T). - return_dict (
bool, por defectoTrue) — Si se debe devolver unStableAudio3DiTModelOutput.
Devuelve
StableAudio3DiTModelOutput o tuple
el campo de velocidad predicho, con la misma forma que hidden_states.
StableAudio3DiTBlock
clase 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 )
Bloque de transformador DiT SA3 individual.
Orden de operaciones:
- Autoatención modulada por AdaLN (RoPE parcial, normalización RMS QK)
- Atención cruzada al contexto de texto/duración (pre-norm RMS simple)
- Red neuronal feed-forward SwiGLU modulada por AdaLN
La modulación AdaLN es to_scale_shift_gate + global_modulation, dividida en seis fragmentos (scale_attn, shift_attn, gate_attn, scale_ff, shift_ff, gate_ff). Cada rama con compuerta se escala mediante sigmoid(1 - gate). La atención cruzada no está modulada por AdaLN, lo que coincide con la referencia (cross_attend_norm es una normalización RMS simple).
Cuando se proporciona local_seq (inpainting), se proyecta por bloque mediante to_local_embed y se suma a las posiciones de audio del flujo residual después de la atención cruzada (y antes de la red feed-forward), coincidiendo con la referencia.
StableAudio3DiTModelOutput
clase diffusers.models.transformers.transformer_stable_audio3.StableAudio3DiTModelOutput
< source >( sample: Tensor )
La salida de StableAudio3DiTModel.