SoundGame/tools/stable-audio-3/optimized/tflite/README.md

sa3_tflite — Stable Audio 3 on CPU via LiteRT / TFLite

Portable CPU inference for Stable Audio 3 — the LiteRT/TFLite sibling of the MLX (Apple Silicon) and TensorRT (NVIDIA GPU) releases. No PyTorch, transformers, or stable-audio-tools at runtime — just ai_edge_litert (LiteRT) driving fully self-contained .tflite graphs through the XNNPACK CPU delegate. Runs anywhere LiteRT runs: macOS / Linux / Windows, x86 / ARM (Windows is x64-only — see Windows).

Quick Install

One line on a fresh machine — installs everything and plays back ~30 seconds of "Impending tribal, epic orchestral buildup":

curl -LsSf https://raw.githubusercontent.com/Stability-AI/stable-audio-3/main/optimized/tflite/bootstrap.sh | bash

Windows — stock PowerShell (no bash/git needed; the curl|bash line above requires Git Bash or WSL on Windows):

irm https://raw.githubusercontent.com/Stability-AI/stable-audio-3/main/optimized/tflite/bootstrap.ps1 | iex

Already cloned the repo? Run from inside optimized/tflite/:

# macOS / Linux / Windows Git Bash
./install.sh                                              # one-time setup
./sa3 --prompt "Impending tribal, epic orchestral buildup" --play           # generates + plays

:: Windows cmd / PowerShell
install.bat
sa3.bat --prompt "Impending tribal, epic orchestral buildup" --play
:: PowerShell users can use .\sa3.ps1 instead of sa3.bat — same flags, and
:: Ctrl-C exits cleanly (cmd's "Terminate batch job (Y/N)?" prompt is a .bat quirk)

Three models, four 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)

Precision variants (--precision)

Same flag as the TensorRT CLI's --precision; the values use the wXaY naming (weights/activations bit-widths). One flag switches the DiT, the codec decoder, and (for audio-to-audio / inpainting) the encoder together; T5Gemma stays single-precision. All variants keep the full feature set: variable length (any --seconds, odd or even latent counts) and variable batch (batched or sequential CFG). Missing files lazy-download from HuggingFace on first use.

--precision	legacy name	size (sm DiT / medium DiT / codecs)	quality	CPU speed
fp32 (default)	—	1.8 GB / 5.8 GB / 0.2–1.8 GB	reference	1× — on CPU this is also the fast choice
w16a32	fp16	0.9 / 2.9 / 0.1–0.9 GB	≈lossless (fp16 weights, fp32 activations; 62–75 dB per-forward)	1.5–3× slower, model-dependent (XNNPACK dequantizes per-matmul)
w8a32	woint8	0.45 / 1.5 / 0.05–0.5 GB	GPTQ int8 weights — codecs transparent (40–46 dB); DiT gives a different but plausible sample	≈fp32
w8a8-dyn	dynint8	0.45 / 1.5 / 0.05–0.5 GB	lowest (int8 weights + activations, per-invoke dynamic scales)	fastest (~1.2–1.3×)

(Names follow the wXaY convention — weight/activation bit-widths, as in LLM releases: quantization here touches the FULLY_CONNECTED weights; activations stay fp32 except w8a8-dyn, whose int8×int8 matmuls are what make it the only faster-than-fp32 variant. All bit-widths are per-channel weight grids; "16" is fp16 — there is no int16 variant. Speed factors measured on an Apple M4 Pro MacBook Pro, XNNPACK, 8 threads.)

Two things worth knowing, both counter-intuitive:

• Quantization buys size, not speed, on CPU. Weight-only int8 dequantizes to fp32 before each matmul, so they run at fp32 speed; fp16 is slower than fp32. Only dynint8 (quantized activations → true int8 matmuls) is faster.
• DiT quantization error compounds. The 8-step sampler is chaotically sensitive: per-step weight error turns into a different (still plausible) sample rather than a noisy version of the fp32 one. Judge DiT precisions by ear; decoder precisions by PSNR (they run once on a fixed latent).

One CFG note: batched and sequential CFG are bit-identical for fp32/w8a32 and inaudibly different (~80 dB) for w16a32, but under w8a8-dyn the batch=2 invoke shares activation-quantization scales across the cond/uncond rows, so batched CFG yields a different plausible sample than sequential. Pass --no-cfg-batched with w8a8-dyn when you need run-to-run reproducibility against sequential baselines.

# quarter-size models, same speed and near-identical quality on the decoder side
./sa3 --prompt "lofi house loop" --dit sm-music --decoder same-s --precision w8a32

# fastest CPU inference (quality tradeoff)
./sa3 --prompt "lofi house loop" --dit sm-music --decoder same-s --precision w8a8-dyn

Encoder variants (used by audio-to-audio / inpainting) exist at the same precisions; int8 weights are GPTQ-calibrated on real audio, like the DiT and codec (measured latent PSNR vs fp32, same-s / same-l): w16a32 66 / 71 dB (≈lossless), w8a32 36 / 46 dB, w8a8-dyn 24 / 30 dB. For the most quality-sensitive inpainting you can still pin --encoder-precision fp32.

--dit-precision / --decoder-precision / --encoder-precision override the shared flag per component — useful because they react to quantization very differently (the codec's precision maps directly to audio fidelity; the DiT's changes which sample you get; the encoder's affects a2a/inpaint latents):

# fastest DiT, reference-quality codec
./sa3 --prompt "lofi house loop" --dit sm-music --decoder same-s \
      --precision fp32 --dit-precision w8a8-dyn

Install

./install.sh

install.sh is uv-based. On a fresh machine it will:

1. Install 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 model file is downloaded from HF on first use and symlinked into models/tflite/ from the HuggingFace cache.

Portable CPU (no GPU required). Python 3.9+. ./install.sh --python 3.12 to pin a different Python.

Windows

The stack runs natively on Windows x64 — no WSL needed. ai-edge-litert ships win_amd64 wheels for Python 3.10–3.13 (x64 only; no Windows-ARM wheels), so install a Python in that range from python.org (check "Add python.exe to PATH"), then from optimized\tflite\:

install.bat                          :: one-time setup (python -m venv + pip)
sa3.bat --prompt "lofi house loop" --dit sm-music --decoder same-s

install.bat / sa3.bat are the Windows twins of ./install.sh / ./sa3 (plain venv + pip; uv not required). Notes:

• ffmpeg is optional — only needed for mp3 / 24-bit / non-44.1 kHz --init-audio inputs: winget install ffmpeg (or choco install ffmpeg).
• Symlinks: downloaded weights are exposed via symlinks where possible; without Developer Mode, Windows disallows creating them, so the downloader automatically falls back to a hardlink (zero-copy) or a plain copy. No action needed either way.
• Long paths: the HuggingFace cache nests deeply — if downloads fail with path-length errors, enable Windows long paths once (admin PowerShell): Set-ItemProperty 'HKLM:\SYSTEM\CurrentControlSet\Control\FileSystem' -Name LongPathsEnabled -Value 1 (the LongPathsEnabled registry switch), then restart the terminal.
• --play uses the stdlib winsound player (macOS uses afplay).

Run

./sa3 is a thin shell wrapper around .venv/bin/python scripts/sa3_tflite.py "$@" that prompts to run ./install.sh if uv or .venv/ isn't set up.

# 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, chunked SAME-L decode)
./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

# Generate + play immediately (afplay/winsound/aplay; 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.

Use --threads to control the XNNPACK CPU thread count (default 8).

Web UI (gradio)

A browser UI over the same TFLite backend — every generation mode plus a spectrogram player, a history panel, and playback options (auto-play, infinite radio with pre-generation, loop, hotswap, MP3/WAV save).

./sa3-gradio                       # sm-music, opens a public share link
./sa3-gradio --dit medium          # start on the medium DiT
./sa3-gradio --no-share            # local-only (http://localhost:7860)
./sa3-gradio --port 7871           # pick the port

./sa3-gradio is a wrapper around .venv/bin/python scripts/sa3_gradio.py; on first run it offers to install the UI-only extras (gradio, pillow, soundfile) into .venv. In the browser:

• Precision dropdown next to the model picker — fp32 (default) / w16a32 / w8a32 / w8a8-dyn, applied to the DiT + codec.
• LoRA panel — upload a .safetensors adapter or pick one from loras/<model>/ (e.g. loras/sa3-medium/), set its strength, stack several; settings are remembered per model. The panel is disabled under a quantized precision (see the LoRA note above) and has no per-step range (frozen graph).
• Audio-to-audio / inpainting — upload a reference clip; inpainting adds a start/end range. Both are combinable.

Generated WAV/MP3s land in output/gradio/. First use of a model/precision loads its weights (cached after).

Without the wrapper

.venv/bin/python scripts/sa3_tflite.py --prompt "..." --dit medium --decoder same-l
# or, after `source .venv/bin/activate`:
python scripts/sa3_tflite.py --prompt "..." --dit medium --decoder same-l

Speed & memory

This is a CPU path — it trades the GPU releases' speed for portability. The small models comfortably beat realtime on a modern laptop CPU; medium is slower (its DiT is ~5.8 GB fp32 and it chunk-decodes SAME-L). Use ≥ ~20 s clips: very short clips have too few latent tokens for the sampler to settle into a coherent loop. Throughput scales with --threads up to your physical core count (4–8 is the usual sweet spot; more threads on a short model adds overhead).

For sub-realtime latency on a supported device, prefer the GPU siblings: MLX on Apple Silicon, TensorRT on NVIDIA.

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 (use ≥ ~20 s)
--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. ≠1 runs cond+uncond each step
--apg	1.0	Adaptive Projected Guidance; only matters when --cfg ≠ 1
--cfg-batched	on	When --cfg ≠ 1, run cond+uncond as one batch=2 invoke on the variable-batch DiT (~7–29% faster on Apple-Silicon AMX). --no-cfg-batched → sequential batch=1 dual-pass. Bit-identical at fp32/w8a32 — see Precision variants for w16a32/w8a8-dyn
--lora	—	A .safetensors LoRA adapter (SA3-native/underfit or PEFT) merged into the DiT, optional strength=S; repeat to stack. Requires --dit-precision fp32 or w16a32. See LoRA
--lora-strength	1.0	Default strength for --lora adapters without their own strength=
--init-audio	—	WAV (any format via ffmpeg) 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
--threads	8	XNNPACK CPU threads (all TFLite models run on CPU)
--free-models	on	Free each model after its last use; --no-free-models keeps them resident
--out / -o	(auto)	Relative → output/<file>; absolute → as-is. 16-bit PCM stereo @ 44.1 kHz, trimmed to exactly --seconds
--play	off	After writing, play the WAV: afplay (macOS) / winsound (Windows) / aplay (Linux); Ctrl-C stops both

All .tflite models are fp32 except T5Gemma, which is fp16 (numerically lossless there). There is no dtype knob: on CPU, int8/fp16 weights buy size, not speed (XNNPACK dequantizes to fp32 to matmul), and int8 costs quality on the DiT — so this release ships the fp32 graphs directly. (See "Notes on the design".)

LoRA

Apply a LoRA finetune by merging it into the DiT's weight buffers at load — same adapters and semantics as the MLX backend's --lora:

# one adapter
./sa3 --dit medium --decoder same-l --prompt "progressive metal" \
      --lora plini-sa3-380.safetensors

# per-adapter strength; stack several
./sa3 --dit medium --decoder same-l --prompt "..." \
      --lora a.safetensors strength=0.8 --lora b.safetensors strength=0.5

The adapter is a .safetensors file (SA3-native train_lora.py / underfit output, or a PEFT adapter directory) — pickle .ckpt/.pt is refused. All nine adapter types (lora, dora-rows/cols, bora, the four -xs) are supported; the base must match --dit. The merge is written into a cached copy of the DiT under models/tflite/lora_cache/ (keyed by adapter + strength content hash) and reused on repeat runs, so the ~5–15 s patch cost is paid once. A medium fp32 cache entry is ~5.4 GB — delete lora_cache/ to reclaim.

Requires --dit-precision fp32 (default) or w16a32. LoRA on the quantized DiTs (w8a32 / w8a8-dyn / w4a32) isn't figured out yet — those store weights as GPTQ-calibrated int8, so merging would mean dequantize → add the LoRA delta → re-quantize, and a naive re-quant throws away the GPTQ error-feedback grid (you'd get a model that's both LoRA-adapted and degraded). Doing it well needs a GPTQ pass over the merged weights; until then --lora refuses quantized precisions (the CLI errors, the web UI disables the LoRA panel). Use fp32/w16a32 for LoRA — on CPU fp32 is the fast-and-accurate choice anyway, so this is rarely a real constraint. Per-step gating (steps=) is also MLX-only — a frozen TFLite graph merges weights once at load; use optimized/mlx for step-gated LoRA.

Files

sa3_tflite/
├── sa3                            ← shell wrapper (use this)
├── install.sh                     ← uv bootstrap (run once)
├── sa3.bat / install.bat          ← Windows twins of sa3 / install.sh
├── bootstrap.sh                   ← one-line curl installer
├── README.md
├── requirements.txt               ← ai_edge_litert, numpy, sentencepiece, soundfile, huggingface_hub
├── output/                        ← default landing zone for generated WAVs
├── scripts/
│   ├── sa3_tflite.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)
└── models/
    ├── tokenizer.model            ← SentencePiece model, BUNDLED (~4 MB; T5Gemma tflite is encoder-only)
    ├── defs/
    │   └── tflite_pipeline.py     ← Tokenizer + T5Gemma front-end + pingpong schedule + sampler + WAV
    └── tflite/                    ← .tflite models (auto-downloaded; ~2.3 GB small, ~9.5 GB medium)
        ├── t5gemma/encoder_fp16.tflite        564 MB   text encoder (fp16)
        ├── sa3-sm-music/dit_fp32.tflite       1.8 GB   small music DiT (conditioner baked in)
        ├── sa3-sm-sfx/dit_fp32.tflite         1.8 GB   small sfx DiT (conditioner baked in)
        ├── sa3-m/dit_fp32.tflite              5.8 GB   medium DiT (conditioner baked in)
        ├── same-s/{enc,dec}_fp32.tflite       ~220 MB each   shared sm-* codec
        └── same-l/{enc,dec}_fp32.tflite       ~1.8 GB each   medium codec

The DiT graphs are baked-I/O: the conditioner (prompt-padding + seconds embedder) and the patch/unpatch are compiled into the graph, so the DiT takes the raw T5Gemma output directly and the decoder emits audio directly. The two small DiTs share the SAME-S codec (bit-exact between checkpoints), so only one set of small-codec files is shipped.

Auto-download from HuggingFace

Model files aren't bundled — they're pulled from stabilityai/stable-audio-3-optimized (under tflite/…) on first use and symlinked into models/tflite/ from the HF cache. No duplication. Anonymous downloads work but are rate-limited; huggingface-cli login with a free read-only token lifts the cap. The SentencePiece tokenizer (models/tokenizer.model, ~4 MB) is the one weight that IS committed, since the .tflite T5Gemma is encoder-only.

Notes on the design

• Baked-I/O varlen graphs. Each .tflite is a single self-contained graph with the conditioner and patch/unpatch in-graph, accepting a variable sequence length — so one file serves any --seconds. The DiT is a 6-input graph (x, t, t5_hidden, t5_mask, seconds, local_add_cond); feed raw T5 outputs and the in-graph conditioner handles prompt-padding + seconds-embedding.
• fp32 everywhere (except fp16 T5Gemma). On CPU, quantizing buys size, not speed — XNNPACK dequantizes int8/fp16 weights to fp32 to matmul, so fp16 is actually slower and int8 gives no speedup. And the DiT will not go int8 at quality: per-step error compounds over the 8 chaotic sampling steps into a different (still plausible) sample, not a degraded one. So this release ships the fp32 graphs directly. T5Gemma fp16 is the sole exception — it's numerically lossless there and halves that file.
• Monotonic audio-to-audio schedule. The pingpong schedule applies the LogSNR shift to the normalized [1→0] grid, then scales by σmax, so audio-to-audio (σmax < 1) stays monotonically decreasing while keeping all N distilled steps. σmax = 1.0 (text-to-audio) is bit-identical to the classic schedule.
• SAME-L chunked decode. The SAME-L decoder's dense sliding-window-attention mask is O(T²), so long clips are decoded in overlap-8 windows of 64 latent tokens (the throughput optimum) and stitched. SAME-S has a narrow receptive field and decodes whole.
• CFG (--cfg ≠ 1) combines a cond and an uncond velocity in denoised space (optional APG). The canonical DiT is variable-batch, so by default cond+uncond run as one batch=2 invoke per step (--cfg-batched) — ~7–29% faster on Apple-Silicon (the AMX matrix unit amortizes the weight loads across both rows; measured on an M4 Pro MacBook Pro). --no-cfg-batched falls back to a sequential batch=1 dual-pass (like the TensorRT release, whose engine is static-batch=1); the two are bit-identical at fp32/w8a32 (~80 dB at w16a32; w8a8-dyn diverges by design — batch-shared activation scales).

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 https://stability.ai/license.