# sa3_mlx — Stable Audio 3 in pure MLX Apple-Silicon-native inference for **Stable Audio 3**, with no PyTorch, transformers, or stable-audio-tools at runtime. ## Quick Install One line on a fresh Apple Silicon Mac — installs everything and plays back ~2 minutes of "Impending tribal, epic orchestral buildup": ```bash curl -LsSf https://raw.githubusercontent.com/Stability-AI/stable-audio-3/main/optimized/mlx/bootstrap.sh | bash ``` Already cloned the repo? Run from inside `optimized/mlx/`: ```bash ./install.sh # one-time setup ./sa3 --prompt "Impending tribal, epic orchestral buildup" --play # generates + plays ``` ## Three models, three modes | `--dit` | model | best for | |------------|--------------------|--------------------------------| | `sm-music` | sa3-sm-music (50 M block) | fast music generation | | `sm-sfx` | sa3-sm-sfx (50 M block) | sound effects | | `medium` | sa3-medium-ARC (1.4 B) | higher-quality music, slower | | mode | flags | example | |------------------|-----------------------------------------------|----------------------------------| | text-to-audio | `--prompt P` | new clip from a description | | audio-to-audio | `--prompt P --init-audio IN.wav --init-noise-level σ` | variation of an existing clip | | inpainting | `--prompt P --init-audio IN.wav --inpaint-range "S,E"` | regenerate one section, keep rest | | CFG + negative | `--cfg 3.0 --negative-prompt P_NEG` | steer toward / away from prompts | ``` prompt ─▶ T5Gemma encoder ─▶ DiT pingpong sampler ─▶ SAME-S/L decoder ─▶ WAV ▲ optional: encoder + init audio (audio-to-audio / inpaint) ``` ## Install ```bash ./install.sh ``` `install.sh` is uv-based. On a fresh machine it will: 1. Install [uv](https://github.com/astral-sh/uv) via the official curl installer if it's missing (prompts y/N; `-y` skips the prompt). 2. Create a project-local `.venv/` with managed Python 3.11. 3. `uv pip install` the runtime deps into the venv (much faster than pip). 4. Ask which DiT bundles to download from HuggingFace (`stabilityai/stable-audio-3-optimized`). Each pick pulls its matching audio codec; T5Gemma (the shared text encoder) is downloaded once. Already-present weights are skipped. End-to-end on a fresh machine: **~10 seconds** + weight downloads. > Don't want to pre-pick bundles? Skip install entirely and just run > `./sa3 --prompt …` — any missing weight file is downloaded from HF on > first use and symlinked into `models/mlx/` from the HuggingFace cache. Apple Silicon only (MLX is Metal-backed). Python 3.10+. `./install.sh --python 3.12` to pin a different Python. ## Run `./sa3` is a thin shell wrapper around `uv run python scripts/sa3_mlx.py "$@"` that prompts to run `./install.sh` if uv or `.venv/` isn't set up. ```bash # Text-to-audio ./sa3 --prompt "lofi house loop" --dit sm-music --decoder same-s --out lofi.wav # Sound effects ./sa3 --prompt "footsteps on gravel" --dit sm-sfx --decoder same-s --out steps.wav # Higher-quality music (medium DiT) ./sa3 --prompt "A beautiful piano arpeggio grows into a cinematic climax" \ --dit medium --decoder same-l --seconds 30 --out piano.wav # Audio-to-audio variation (σmax 0.4-0.8 typical) ./sa3 --prompt "jazz fusion with electric piano" --dit sm-music --decoder same-s \ --init-audio funk.wav --init-noise-level 0.7 --out funk_jazz.wav # Inpaint seconds 4-7 ./sa3 --prompt "explosive drum break" --dit sm-music --decoder same-s \ --init-audio funk.wav --inpaint-range "4,7" --out funk_drums.wav # CFG + negative prompt ./sa3 --prompt "ambient drone" --cfg 3.0 --negative-prompt "drums, vocals" \ --dit sm-music --decoder same-s --out drone.wav # Apply a LoRA finetune (merged into the DiT at load; base must match --dit) ./sa3 --prompt "arabic maqam oud taqsim" --dit medium --decoder same-l \ --lora ./my_lora.safetensors --lora-strength 1.0 --out maqam.wav # Per-LoRA strength + sampling-step gating (skipping step 1 often helps). # steps= is 1-based inclusive: 2-8, 2- (2..last), -4 (1..4), 3 (just step 3). # Step-gated adapters are re-merged in place at the few step boundaries # (~80 ms each on medium) — no extra memory, every step runs at full speed. ./sa3 --prompt "progressive metal instrumental" --dit medium --decoder same-l \ --lora plini.safetensors strength=0.8 steps=2- \ --lora rain_texture.safetensors steps=-4 --out prog.wav # Generate + play immediately (afplay; Ctrl-C stops both) ./sa3 --prompt "rainforest" --dit sm-sfx --decoder same-s --play # All options + categorised examples ./sa3 --help ``` Omit `--dit` / `--decoder` for an interactive arrow-key picker. Omit `--prompt` for a stdin prompt. Relative `--out` paths land in `output/` (auto-created); absolute paths are honoured as-is. The output path is printed prominently as a `▸ saved` line at the end of each run. ### Without the wrapper ```bash uv run python scripts/sa3_mlx.py --prompt "..." --dit medium --decoder same-l # or, after `source .venv/bin/activate`: python scripts/sa3_mlx.py --prompt "..." --dit medium --decoder same-l ``` ## Web UI (gradio) ```bash ./sa3-gradio # public gradio.live share link by default ./sa3-gradio --no-share # local-only (http://127.0.0.1:7860) ./sa3-gradio --dit medium # initial model (switchable in the UI) ``` Every generation mode is wired: text-to-audio, CFG 0–10 (0 = the negative prompt takes over, 0.5 = halfway between both prompts, >1 = extrapolate), audio-to-audio (guide audio + σmax slider), and inpainting (separate reference audio + start/end range sliders) — a2a and inpainting combine, so an inpainted span can regenerate from the guide audio instead of noise. Each clip renders with a 3-band tinted stereo mel spectrogram (bass=red / mid=green / high=blue). Models hot-swap from the dropdowns and cache in unified memory, so switching back is instant; WAVs land in `output/gradio/`. The UI needs two extra packages (`gradio`, `pillow`) — `./sa3-gradio` offers to install them into `.venv` on first run. ## Speed & memory Measured on **M1 8 GB** (the slowest M-chip). Newer chips are faster — roughly **1.2–1.4× on M2**, **1.5–2× on M3**, **2–3× on M4**, especially on the memory-bandwidth-bound seq=1 attention path. | `--dit` | Wall (10 s clip) | × realtime | Peak RAM | |-----------------------|------------------|------------|----------| | `sm-music` / `sm-sfx` | ~1 s | ~10× | 1.6 GB | | `medium` | ~5 s | ~2× | 3.8 GB | Peak RAM assumes `--free-models` (default on) — T5Gemma is freed after conditioning and the DiT after sampling, so the decoder never competes with the upstream models for memory. 8 GB M1 is fine for everything; disable with `--no-free-models` only if you have headroom and plan to batch multiple runs in one process. ### Benchmark yourself Reproduce the table above on your own machine — sweeps both models across 4 clip lengths (5 s / 30 s / 120 s / 380 s) and prints an ASCII summary. Each cell runs in its own subprocess so peak RAM is measured cleanly per run; weights are pre-warmed via `ensure_local()` so HF download time doesn't pollute the timings. Takes ~3-10 min depending on your chip. ```bash uv run --no-project python scripts/benchmark.py ``` Sample run on **M4 Pro / 48 GB**: ``` ┌───────────┬─────────┬─────────┬───────────┬───────────┬────────────┐ │ model │ decoder │ seconds │ wall (s) │ ×realtime │ peak RAM │ ├───────────┼─────────┼─────────┼───────────┼───────────┼────────────┤ │ sm-music │ same-s │ 5 │ 0.80 │ 6.27× │ 1.62 GB │ │ sm-music │ same-s │ 30 │ 1.53 │ 19.61× │ 1.94 GB │ │ sm-music │ same-s │ 120 │ 4.12 │ 29.13× │ 2.38 GB │ │ sm-music │ same-s │ 380 │ 13.46 │ 28.24× │ 2.58 GB │ ├───────────┼─────────┼─────────┼───────────┼───────────┼────────────┤ │ medium │ same-l │ 5 │ 2.20 │ 2.27× │ 3.82 GB │ │ medium │ same-l │ 30 │ 5.06 │ 5.92× │ 3.89 GB │ │ medium │ same-l │ 120 │ 14.68 │ 8.17× │ 5.21 GB │ │ medium │ same-l │ 380 │ 47.78 │ 7.95× │ 5.05 GB │ └───────────┴─────────┴─────────┴───────────┴───────────┴────────────┘ ``` ## LoRA training Finetune SA3 on your own audio, entirely in MLX (no PyTorch) — a full training CLI that replicates [underfit](https://github.com/dada-bots/underfit)'s conventions, defaults, and checkpoint format (so it can run as underfit's Apple-Silicon backend). Three steps: **pre-encode → train → generate**. See `TRAINING_CONVENTIONS.md` for the complete convention inventory and the torch(MPS)-vs-MLX forward+backward parity results. > **Want a UI instead of the CLI? Use [underfit](https://github.com/dada-bots/underfit).** > It's a full LoRA-training dashboard that drives *this* MLX trainer as its > Apple-Silicon backend — dataset scanning/encoding, live loss + loss-by-timestep > curves, demo MP3s with spectrograms, per-checkpoint gradio launch, and > one-click model downloads. On a Mac it's the recommended way to train; follow > underfit's *Apple-Silicon quickstart*. The commands below are the lower-level > path — good for scripting, CI, or headless runs. **1. Pre-encode** your audio to SAME latents (once, offline — torch-free): ```bash uv run python scripts/pre_encode_mlx.py \ --audio-dir ~/my-clips --output-dir ~/my-latents --codec same-s ``` Each file becomes `.npy` latents `[D, T]` + a `.json` sidecar (duration, padding mask, tags) — the exact format underfit's dataset uses. Use `same-s` for sm-music/sm-sfx, `same-l` for medium. **2. Train.** Training uses the **BASE** checkpoint (`stabilityai/stable-audio-3-*-base`, rectified_flow), *not* the shipped ARC weights inference uses. The base-model npz is **auto-downloaded** from HuggingFace on first run — no flag needed (override with `--dit-weights `; regenerate one yourself from a `-base` repo with `scripts/export_base_npz.py`): ```bash uv run python scripts/lora_train_mlx.py \ --dit sm-music --latents-dir ~/my-latents --lr 1e-4 --name my-lora \ --adapter-type dora-rows --rank 16 --max-steps 2000 ``` Defaults mirror underfit (dora-rows / rank 16, AdamW, uniform sampler + the "full" distribution shift, CFG-dropout 0.1, signal-only masked rectified-flow loss, checkpoint every 1000 steps). To match the full dashboard template config, add `--timestep-sampler trunc_logit_normal --use-effective-length --beta2 0.95 --weight-decay 0.01 --lr-scheduler inverse --lr-warmup 0.995`. Checkpoints land at `output/runs///checkpoints/-step=S-epoch=E.safetensors` with `lora_config` metadata; resume with `--lora-ckpt-path `. Optional training-time demos (RF-Euler inference → mp3, underfit's on-disk format): `--demo-every 1000 --demo-config demos.json`, where `demos.json` is a list of `{"prompt", "cfg", "seed", "steps", "lora_strength"?, "lora_interval_max"?}`. **3. Generate** with your adapter — the checkpoint applies to the shipped ARC model at inference via `--lora` (see "Apply a LoRA finetune" above; also `scripts/sa3_gradio.py --lora ` to open the web UI with it preloaded): ```bash ./sa3 --dit sm-music --decoder same-s --prompt "..." --out finetuned.wav \ --lora "output/runs/my-lora/"*"/checkpoints/my-lora-step=2000-epoch="*.safetensors ``` ### Building your own loop The CLI is assembled from reusable pieces: `inject_trainable_lora` / `save_lora_checkpoint` / `apply_lora_checkpoint` (`models/defs/lora.py` — all 9 LoRA/DoRA/BoRA/-XS types with the no-weight-materialization forward), the RF loss + samplers + distribution shift (`models/defs/training.py`), the pre-encoded dataset + prompt templates (`models/defs/latent_dataset.py`), and `encode_audio` (`models/defs/audio_encoding.py`). Checkpoints use the same safetensors keys + `lora_config` metadata as the PyTorch trainer, so adapters are interchangeable in either direction. ## Flag reference | Flag | Default | Notes | |-----------------------|----------|-----------------------------------------------------------------------| | `--prompt` | (asks) | Text prompt; empty string = unconditional | | `--negative-prompt` | — | CFG uncond branch; only used when `--cfg ≠ 1.0` | | `--dit` | (asks) | `sm-music`, `sm-sfx`, or `medium` | | `--decoder` | (asks) | `same-s` (pairs with sm-*) or `same-l` (pairs with medium) | | `--seconds` | 30 | Output length | | `--steps` | 8 | Pingpong sampler steps; 1 = single forward (fastest), 8 = sweet spot | | `--seed` | random | Set for reproducibility; the chosen seed is printed at the end | | `--cfg` | 1.0 | Guidance scale; 1.0 = off, >1 toward prompt, <1 toward uncond | | `--apg` | 1.0 | Adaptive Projected Guidance; only matters when `--cfg ≠ 1` | | `--init-audio` | — | WAV (44.1 kHz, 16-bit PCM) input for audio-to-audio / inpaint | | `--init-noise-level` | 1.0 | σmax; 0.4–0.8 typical for variation, 1.0 = full regen, >1 = overshoot | | `--inpaint-range` | — | `START,END` seconds; regenerate that span, keep the rest | | `--dit-dtype` | fp16 | DiT compute dtype (decoder always FP32; T5Gemma always fp16) | | `--lora` | — | A `.safetensors` LoRA adapter (SA3-native/underfit or PEFT) with optional `strength=S` and `steps=MIN-MAX` tokens; repeat the flag to stack adapters. Full-range adapters merge at load; step-gated ones re-merge in place at step boundaries. Pickle `.ckpt/.pt` is refused. Base must match `--dit` | | `--lora-strength` | 1.0 | Default strength for adapters without their own `strength=`; 0 = bit-exact bypass, >1 amplifies | | `--free-models` | on | Progressive model freeing; `--no-free-models` keeps them resident | | `--out` | out.wav | Relative → `output/`; absolute → as-is. 16-bit PCM stereo @ 44.1 kHz, trimmed to exactly `--seconds` | | `--play` | off | After writing, play via `afplay`; Ctrl-C stops both processes | ## Files ``` sa3_mlx/ ├── sa3 ← shell wrapper (use this) ├── install.sh ← uv bootstrap (run once) ├── README.md ├── TRAINING_CONVENTIONS.md ← LoRA-training conventions + MPS-vs-MLX parity ├── requirements.txt ├── output/ ← default landing zone for generated WAVs ├── scripts/ │ ├── sa3_mlx.py ← orchestrator CLI (invoked by ./sa3) │ ├── weights.py ← weights manifest + HF auto-download │ ├── examples.py ← shared examples block (--help + post-install) │ ├── install.py ← install.sh's Python half (bundle picker) │ ├── test_all_configs.py ← npz + CLI config sanity tests │ ├── benchmark.py ← wall-time + peak-RAM matrix across model × duration │ ├── lora_train_mlx.py ← LoRA training CLI (underfit conventions) │ ├── pre_encode_mlx.py ← audio → SAME-latent pre-encode (torch-free) │ ├── sa3_gradio.py ← web UI (invoked by ./sa3-gradio; --lora preload) │ ├── test_lora_merge.py ← adapter-math tests (weight-free; run in CI) │ └── parity_forward_{torch,mlx}.py ← torch(MPS)-vs-MLX fwd+bwd parity harness └── models/ ├── defs/ │ ├── sa3_pipeline.py ← sampler + conditioner + unpatch │ ├── t5gemma_mlx.py ← T5Gemma encoder + SentencePiece wrapper │ ├── dit_mlx.py ← small DiT (sm-music + sm-sfx) │ ├── dit_mlx_medium.py ← medium DiT (differential attention) │ ├── same_s_{encoder,decoder}.py ← small codec │ ├── same_l_{encoder,decoder}.py ← large codec │ ├── lora.py ← trainable LoRA/DoRA/BoRA adapters (9 types) │ ├── lora_merge.py ← inference-time LoRA merge + per-step gating │ ├── training.py ← RF loss + timestep samplers + distribution shift │ ├── latent_dataset.py ← pre-encoded dataset + underfit prompt templates │ ├── audio_encoding.py ← waveform → SAME latents (pre-encode bridge) │ └── demo_mlx.py ← training-time RF-Euler demos → mp3 └── mlx/ ← .npz weights (auto-downloaded; ~8.4 GB total) ├── t5gemma_f16.npz 541 MB text encoder + tokenizer ├── dit_sm-music_f16.npz 877 MB DiT + conditioner baked in ├── dit_sm-sfx_f16.npz 877 MB DiT + conditioner baked in ├── dit_medium_f16.npz 2.77 GB DiT + conditioner baked in ├── same_s_{encoder,decoder}_f32.npz ~210 MB each shared sm-* codec └── same_l_{encoder,decoder}_f32.npz ~1.7 GB each medium codec ``` Each DiT `.npz` bundles its small conditioner under a `cond.*` key prefix (learned `padding_embedding` + a tiny SecondsTotalEmbedder), so the orchestrator only loads one file per model. The two small DiTs share the same SAME-S codec weights — they're bit-exact between checkpoints, so only one set of small-codec npz files is shipped. ## Notes on the design - **Mixed precision**: DiT runs FP16 (validated transparent at ~50–57 dB PSNR vs FP32), the audio decoder always runs FP32 (SAME-S's differential attention catastrophically cancels in FP16), T5Gemma always runs FP16. `--dit-dtype fp32` reverts to all-FP32 for bit-exact reproducibility. - **Padding embedding**: empty / short prompts don't see raw zeros for pad positions — they see a learned 768-dim `cond.padding_embedding` baked into each DiT npz. This is what makes "" (empty prompt) produce sensible unconditional audio instead of noise. - **Auto-download**: weights aren't bundled — they're pulled from `stabilityai/stable-audio-3-optimized` on HF on first use and symlinked into `models/mlx/` from the HF cache. No duplication. - **Chunked decoding**: SAME-L can't fit a long clip into one decode pass on 8 GB; the decoder splits into overlap-2 (SAME-S) or overlap-8 (SAME-L) chunks with a uniform-kernel edge treatment that stays bit-exact vs an un-chunked reference at the natural sizes. ## License & attribution Model weights derived from Stability AI's Stable Audio 3 checkpoints. T5Gemma text encoder from Google. Use of the Stable Audio 3 weights is governed by the **Stability AI Community License**. Please refer to the full terms at .