tflite/README.md

# 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

```python
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).