Hugging Face's logo

Diffusers documentation

StableAudio3DiTModel

Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

StableAudio3DiTModel

Diffusion Transformer (DiT) для предсказания скорости rectified-flow, предназначенный для генерации аудио и используемый в Stable Audio 3.

Каждый StableAudio3DiTBlock выполняет:

  1. Self-attention — дифференциальное многоголовое внимание с вращательными позиционными эмбеддингами (RoPE).
  2. Кросс-внимание — обращается к последовательности токенов из текстового энкодера T5Gemma.
  3. Feed-forward — проекция SwiGLU.

Модель обусловлена timestep (экспоненциальные признаки Фурье → линейная проекция) и глобальным вектором обусловливания (эмбеддинг длительности из StableAudio3DurationEmbedder).

StableAudio3DiTModel

класс 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 )

Параметры

  • io_channels (int, по умолчанию 256) — Количество латентных каналов.
  • patch_size (int, по умолчанию 1) — размер временного патча, применяемого перед трансформером.
  • embed_dim (int, по умолчанию 1536) — скрытое измерение Transformer.
  • depth (int, по умолчанию 24) — количество слоёв StableAudio3DiTBlock.
  • num_heads (int, по умолчанию 24) — Количество головок внимания.
  • cond_token_dim (int, по умолчанию 768) — Размерность токенов контекста для кросс-внимания.
  • global_cond_dim (int, по умолчанию 768) — Размерность глобального эмбеддинга длительности.
  • local_add_cond_dim (int, по умолчанию 257) — Количество каналов локально-аддитивного (inpaint) тензора.
  • timestep_features_dim (int, по умолчанию 256) — Выходная размерность признаков временных шагов Фурье.
  • ff_mult (float, по умолчанию 4.0) — коэффициент расширения SwiGLU для полносвязного слоя.
  • num_memory_tokens (int, по умолчанию 64) — Количество обучаемых токенов памяти.
  • use_differential_attention (bool, по умолчанию True) — Включает дифференциальное само- и кросс-внимание.

Базовая архитектура Diffusion Transformer (DiT) в Stable Audio 3.

Модель принимает пакет зашумленных аудиолатентов, скалярный временной шаг, контекст кросс-внимания (проецированные токены текста и длительности) и глобальное эмбеддинг длительности, и предсказывает поле скоростей (цель rectified-flow).

Условие:

  • Контекст кросс-внимания (encoder_hidden_states) проецируется с помощью to_cond_embed.
  • Глобальное эмбеддинг длительности (global_hidden_states) проецируется через to_global_embed, суммируется с эмбеддингом временного шага, а затем расширяется через global_cond_embedder в модуляцию AdaLN для каждого блока.
  • local_add_cond (inpainting) проецируется по блокам с помощью to_local_embed.

num_memory_tokens обучаемых токенов добавляются в начало аудиопоследовательности внутри трансформера и удаляются перед выходной проекцией.

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

Параметры

  • hidden_states (torch.Tensor) — Шумные латентные аудио (batch, io_channels, T).
  • timestep (torch.Tensor) — Шаг диффузии (batch,) в диапазоне [0, 1].
  • encoder_hidden_states (torch.Tensor) — Контекст кросс-внимания (batch, T_ctx, cond_token_dim).
  • global_hidden_states (torch.Tensor) — Глобальное эмбедирование длительности (batch, global_cond_dim).
  • encoder_attention_mask (torch.Tensor, необязательный) — Bool-маска (batch, T_ctx), True = валидный.
  • local_add_cond (torch.Tensor, optional) — Локально-аддитивное (inpaint) условие (batch, local_add_cond_dim, T).
  • return_dict (bool, по умолчанию True) — возвращать ли StableAudio3DiTModelOutput.

Возвращает

StableAudio3DiTModelOutput или tuple

предсказанное поле скоростей, с той же формой, что и hidden_states.

StableAudio3DiTBlock

класс 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 )

Один блок трансформера SA3 DiT.

Порядок операций:

  1. Само-внимание с модуляцией AdaLN (частичный RoPE, нормализация RMS для QK)
  2. Кросс-внимание к контексту текста/длительности (обычный RMS pre-norm)
  3. SwiGLU-функция с модуляцией AdaLN

Модуляция AdaLN представляет собой to_scale_shift_gate + global_modulation, разбитую на шесть частей (scale_attn, shift_attn, gate_attn, scale_ff, shift_ff, gate_ff). Каждая ветвь с гейтом масштабируется множителем sigmoid(1 - gate). Кросс-внимание не модулируется через AdaLN, что соответствует эталонной реализации (cross_attend_norm — это обычный RMS-нормализатор).

Если задан local_seq (inpainting), он проецируется по блокам с помощью to_local_embed и добавляется к позициям аудио в остаточном потоке после cross-attention (и перед feed-forward), что соответствует эталонной реализации.

StableAudio3DiTModelOutput

класс diffusers.models.transformers.transformer_stable_audio3.StableAudio3DiTModelOutput

< >

( sample: Tensor )

Параметры

  • sample (torch.Tensor) — Предсказанное поле скоростей, имеющее ту же форму, что и входные hidden_states.

Вывод StableAudio3DiTModel.

Update on GitHub

←SkyReelsV2Transformer3DModel StableAudioDiTModel→