feat: add Stable Audio 3 text-to-audio model

[shreyaskarnik]

Summary

• Adds Stable Audio 3 (SA3) diffusion-based text-to-audio generation support
• Architecture: DiT backbone (20L, 1024d, AdaLN) + SAME autoencoder (4096x downsampling) + T5Gemma text encoder (768d)
• Generates stereo 44.1kHz audio from text prompts — sound effects and music up to 30s
• Supports Euler and ping-pong samplers with classifier-free guidance
• Also adds model_config.json fallback to load_config() (SA3's HF repo uses this instead of config.json)
• Also adds --duration, --sampler, and --seed CLI args to generate.py

Converted weights

Uploaded to HuggingFace:

• mlx-community/stable-audio-3-small-sfx — sound effects
• mlx-community/stable-audio-3-small-music — music generation

Usage

# Sound effects
python -m mlx_audio.tts.generate \
  --model mlx-community/stable-audio-3-small-sfx \
  --text "Dog barking next to a waterfall" \
  --duration 10 --ddpm_steps 8 --verbose

# Music
python -m mlx_audio.tts.generate \
  --model mlx-community/stable-audio-3-small-music \
  --text "Lo-fi hip hop beat" \
  --duration 30 --ddpm_steps 8 --verbose

Python API

from mlx_audio.tts.utils import load
from mlx_audio.tts.generate import generate_audio

model = load("mlx-community/stable-audio-3-small-sfx")
generate_audio(
    text="Thunder rolling in the distance",
    model=model,
    duration=10.0,
    ddpm_steps=8,
    output_path="output/",
    verbose=True,
)

Architecture notes

• DiT: 20-layer transformer with exponential Fourier features for timestep embedding, RoPE (partial rotation on first 32/64 dims), memory tokens, and cross-attention for text conditioning
• SAME decoder: Transformer-based upsampler with differential attention, DynamicTanh (DYT) normalization, chunk_midpoint_shift for boundary artifacts, and SoftNormBottleneck
• T5Gemma: Text encoder runs via PyTorch/transformers (lazy-loaded on first generate call). This is the one PyTorch dependency at inference time — no MLX port of T5Gemma exists yet
• Samplers: Euler (ODE) and Ping-pong (re-noising) with distribution shift schedule
• sanitize(): Handles weight normalization merging (weight_g + weight_v → weight), Conv1d transposition (PyTorch → MLX layout), and is idempotent for already-converted weights

Test plan

• Convert SFX model from HF via mlx_audio.convert pipeline
• Convert Music model from HF via mlx_audio.convert pipeline
• Generate audio via CLI (python -m mlx_audio.tts.generate)
• Generate audio via Python API (load() + generate_audio())
• Load from local path
• Load from HuggingFace hub (mlx-community/stable-audio-3-small-sfx)
• Verify model.sanitize() is idempotent (runs on already-converted weights without corruption)
• Strict weight loading passes (no missing/extra parameters)

[lucasnewman]

@shreyaskarnik What's the status here? Are you still working on this?