Stable Audio 3 — LiteRT / TFLite models (variable-length, fp32, self-contained)
Portable CPU (and browser-via-WASM) LiteRT builds of the SA3 pipeline, parallel to
the mlx/, onnx/, and TensorRT releases. Runtime: ai_edge_litert (LiteRT); the
XNNPACK delegate accelerates all of these.
Every model is variable-length — a single graph runs any sequence length via
interpreter.resize_tensor_input(...) — except T5Gemma (fixed 256 text tokens, which
is correct: it encodes the prompt, independent of audio duration). The pre/post glue
(conditioning, patch/unpatch) is baked into the graphs, matching the onnx/ tree,
so these are drop-in self-contained (no host-side conditioner/patch code needed).
All fp32 except T5Gemma (fp16 — its bf16 source makes fp16 lossless).
Files
| path | component | precision | I/O (dynamic axis = L latent tokens; N = L·4096 audio samples) |
|---|---|---|---|
sa3-sm-music/dit_fp32.tflite |
DiT small (sa3-sm-music) | fp32 | x[1,256,L], t[1], t5_hidden[1,256,768], t5_mask[1,256], seconds_total[1], local_add_cond[1,257,L] → velocity[1,256,L] |
sa3-m/dit_fp32.tflite |
DiT medium (sa3-medium) | fp32 | (same 6-in → velocity) |
sa3-sm-sfx/dit_fp32.tflite |
DiT small-SFX (sa3-sm-sfx) | fp32 | (same 6-in → velocity) |
same-s/enc_fp32.tflite |
SAME-S encoder | fp32 | audio [1,2,N] → latents [1,256,L] |
same-s/dec_fp32.tflite |
SAME-S decoder | fp32 | latents [1,256,L] → audio [1,2,N] |
same-l/enc_fp32.tflite |
SAME-L encoder | fp32 | audio [1,2,N] → latents [1,256,L] |
same-l/dec_fp32.tflite |
SAME-L decoder | fp32 | latents [1,256,L] → audio [1,2,N] |
t5gemma/encoder_fp16.tflite |
T5Gemma text encoder | fp16 | input_ids[1,256] i32, attention_mask[1,256] i32 → [1,256,768] (fixed 256) |
Conditioning is baked into the DiT (prompt padding + seconds embed + concat), so the
DiT takes the raw T5Gemma output directly — I/O identical to onnx/…/dit.onnx.
local_add_cond is a zeros input [1,257,L] (resize its dim-2 to L along with x).
Sizes: small/sfx DiT 1.7 GB · medium DiT 5.4 GB · SAME-S enc/dec 0.2 GB each · SAME-L enc/dec 1.7 GB each · T5Gemma 0.5 GB. Total ≈ 13 GB.
Variable-length usage
from ai_edge_litert import interpreter as tfl
it = tfl.Interpreter(model_path="same-s/dec_fp32.tflite", num_threads=4)
xi = it.get_input_details()[0]["index"] # latents [1,256,L]
it.resize_tensor_input(xi, [1, 256, L]); it.allocate_tensors()
it.set_tensor(xi, latents); it.invoke()
audio = it.get_tensor(it.get_output_details()[0]["index"]) # [1,2,L*4096]
The graphs recompute RoPE / attention masks in-graph from the live length, so XNNPACK stays fully delegated and per-length latency matches the equivalent fixed-shape export — no dynamic-shape penalty.
Per-model notes (all validated lossless vs the fp32 / MLX reference)
- DiT (small/medium/sfx): >100 dB vs both the torch reference and the ONNX-ORT refs, at every length 16 → 4096 latent tokens (~1.5 s → 6:20), incl. odd lengths. Whole-sequence, unmasked attention → ~linear cost. Latents are softnorm (std≈1).
- SAME-S enc/dec: ~linear, run whole. ~89–98 dB vs reference.
- SAME-L enc/dec: dense sliding-window-attention mask → cost is O(L²) on whole-sequence runs. For long clips, chunk with overlap = 8 latent tokens (accuracy is constant for any chunk ≥ the window; ~64 latent tokens is the throughput sweet spot). 98–112 dB vs reference. SAME-S needs only overlap = 2.
- T5Gemma: fixed 256 by design (prompt, padded+masked; independent of audio duration). Matches the MLX (padded) and TRT (static) paths.
Running the full pipeline
prompt → T5Gemma → DiT (8-step rectified-flow) → SAME decoder → WAV (conditioning
and wave-conversion are now in-graph). A reference CLI (sa3_tflite.py in the
speed-metal repo) drives it and prints per-stage timing. The only host-side pieces
left are the SentencePiece tokenizer (prompt → input_ids) and the 8-step sampler loop.
Encoders run the reverse (reference audio → latents).