README.md

---
license: apache-2.0
library_name: litert
pipeline_tag: text-to-speech
tags:
  - text-to-speech
  - tts
  - dialogue
  - rq-transformer
  - moshi
  - mimi
  - dia2
  - litert
  - tflite
  - on-device
base_model: nari-labs/Dia2-1B
---

# Dia2-1B — LiteRT (on-device, two-speaker dialogue TTS)

[Dia2-1B](https://huggingface.co/nari-labs/Dia2-1B) (Nari Labs, Apache-2.0) — a two-speaker
**dialogue** text-to-speech model — re-authored to run on-device with LiteRT `CompiledModel`.
Verified on a Pixel 8a (Tensor G3).

Dia2 is a Moshi-style **RQ-Transformer**. Once per 12.5 Hz frame a 30-layer *temporal* transformer
emits a word-timing action plus Mimi codebook 0; a 3-layer *depformer* then autoregressively fills
the remaining 31 codebooks for that same frame. [Mimi](https://huggingface.co/kyutai/mimi) (32
quantizers) decodes the codes to 24 kHz audio.

Everything runs on **CPU (fp32)**, because fp16 collapses these deep stacks on ARM XNNPACK. The
KV caches, RoPE, embedding sums, the depformer's projections and all sampling live on the host; the
graphs are pure step functions.

> **Correction (2026-07-10).** An earlier version of this card said the Mali ML Drift GPU delegate
> rejects the language models' KV-step `FULLY_CONNECTED` weight shapes. That rejection is real on
> LiteRT **2.1.3** and **fixed in 2.1.5**. The depformer's own compile failure was in *our* graph: a
> rank-5 reshape inside the fused-QKV authoring (ML Drift's maximum tensor rank is 4). Slicing the
> last dimension into thirds gives **237/237 nodes delegated**; it then miscomputed at both default
> and FP32 precision, which is the known BMM + broadcast-`ADD` bug, and pre-expanding the attention
> mask host-side from `[1,1,1,D]` to `[1,NH,1,D]` brings it to **corr 1.000000**.
>
> **Measured end to end on a Pixel 8a:** with the depformer on the GPU the audio is bit-identical to
> the CPU path (4288/4288 codebook tokens equal, waveform corr 1.000000, RMS diff 0.000000), but it
> is **no faster** — 3906 depformer calls cost 76.4 s on CPU and 78.3 s on GPU. `CompiledModel.run()`
> only *enqueues*; the GPU work is paid inside the first `readFloat()` (17.1 s enqueue + 61.2 s
> readback), so timing `run()` alone suggests a 4.4x win that does not exist. Per call the GPU costs
> 21.1 ms against the CPU's 19.7 ms: a 3-layer single-token step graph cannot amortise dispatch and
> synchronisation. The GPU is not the obstacle, and it is also not the answer at this graph size.

## Files

| File | Size | Input | Output |
|---|---|---|---|
| `dia2_temporal_fp32.tflite` | 3.0 GB | emb `[1,1,1024]`, RoPE cos/sin `[1,1,1,128]`, mask `[1,1,1,257]`, packed K/V `[1,240,256,128]` | hidden `[1,1,1024]`, action `[1,1,2]`, cb0 `[1,1,2050]`, new K/V |
| `dia2_depformer_wi{0,1,2}_fp32.tflite` | 164 MB each | dep_in `[1,1,1024]`, RoPE, mask `[1,1,1,32]`, packed K/V `[1,24,31,128]` | hidden `[1,1,1024]`, new K/V |
| `dia2_mimi_dequant.tflite` | 68 MB | codes `[1,32,1]` float | latent `[1,512,1]` |
| `dia2_mimi_decode_t256.tflite` | 164 MB | latent `[1,512,256]` | audio `[1,1,491520]` @ 24 kHz |
| `dia2_combined_main.f16` / `dia2_combined_second.f16` | 101 MB each | — | text embedding × projection, `[49280,1024]` |
| `dia2_temporal_audio.f16` | 134 MB | — | audio embeddings, `[32,2050,1024]` |
| `dia2_dep_audio.f16` / `dia2_dep_logits.f16` | 130 MB each | — | depformer embeddings / logit weights, `[31,2050,1024]` |
| `dia2_dep_in.f16` | 6 MB | — | depformer input projections, `[3,1024,1024]` |
| `dia2_constants.json` | — | — | delays, weight schedule, token ids |
| `dia2_prefix.json` | 13 kB | — | baked two-speaker voice prompt |

## Three things that are easy to get wrong

**1. Both text streams carry real word tokens.** Channels 0 and 1 are *not* new-word/pad markers.
On a new word the main stream emits the word's first text token while the second stream emits
`NEW_WORD`; during the padding frames that follow, the main stream drains the rest of the word and
the second stream drains a two-word lookahead. Feeding markers instead produces fluent, confident,
completely wrong speech.

**2. Undo the delay pattern before decoding.** Codebook `cb` lags the aligned timeline by 16 frames
(cb0) or 18 frames (the rest). Codes are stored at `audio[cb][t+1]`, so
`aligned[cb][τ] = audio[cb][delay[cb] + τ]`, and the result is `max(delay)` frames shorter. Skipping
this yields muffled, unintelligible audio.

**3. Mimi decode must be a single pass.** The decode path is upsample → **causal** decoder
transformer → SEANet, so its receptive field is unbounded. Decoding disjoint windows starts each one
with no history and costs ~13% relative error (corr 0.991 against a full-sequence decode). The graph
here spans 256 frames; leave the unused tail zeroed and causality makes it exact — corr **0.999999**.

## The speaker is sampled

With no voice prefix Dia2 **samples the speaker identity**, so the voice changes on every run (median
F0 wanders over a ~120 Hz range). Classifier-free guidance does *not* fix this — measured over 8
matched seeds, the F0 spread is 144 Hz at `cfg_scale=1.0` and 134 Hz at `2.0`; what guidance buys is
steadier output levels. The model's own remedy is a **voice prefix**.

Building a prefix normally needs Whisper word timings and a Mimi *encoder*, both host-only, so
`dia2_prefix.json` ships a **precomputed** prompt (aligned Mimi codes, `new_word_steps`, prefix word
entries). On device only the warm-up runs: the temporal transformer replays the prompt to prime both
KV caches — no Mimi encoder, no sampling, no depformer. The generated speakers then track their
prompts (measured on device: S1 214 Hz / S2 114 Hz, against prompts of 247 Hz / 88 Hz; without a
prefix S2 never drops below 214 Hz).

Classifier-free guidance (`cfg_scale = 2.0`, Dia2's default) is subtle: the guided logits
`uncond + scale·(cond − uncond)` only **select** the top-k candidate set, while the draw is a
temperature softmax over the **conditional** logits restricted to that set. It therefore needs a
second, unconditional branch (text forced to `zero`/`pad`, same audio codes, its own KV cache).

## Usage — Kotlin (Android, LiteRT `CompiledModel`)

```kotlin
import com.google.ai.edge.litert.Accelerator
import com.google.ai.edge.litert.CompiledModel

val options = CompiledModel.Options(Accelerator.CPU)
val temporal = CompiledModel.create("$dir/dia2_temporal_fp32.tflite", options, null)
val depformer = Array(3) { CompiledModel.create("$dir/dia2_depformer_wi${it}_fp32.tflite", options, null) }

// One temporal step: embedding + RoPE + additive mask + packed KV cache.
val inputs = temporal.createInputBuffers()
val outputs = temporal.createOutputBuffers()
inputs[0].writeFloat(embedding)       // [1,1,1024] = text embed + 32 audio embeds, summed on host
inputs[1].writeFloat(ropeCos(frame))  // [1,1,1,128]
inputs[2].writeFloat(ropeSin(frame))
inputs[3].writeFloat(cache.mask())    // [1,1,1,257]; -3e38 on unwritten slots, 0 on the tail
inputs[4].writeFloat(cache.keys)      // [1,240,256,128]
inputs[5].writeFloat(cache.values)
temporal.run(inputs, outputs)
// outputs: hidden [1024], action [2], cb0 logits [2050], new K/V (append to the cache)
```

## Usage — Python (LiteRT `CompiledModel`)

```python
import numpy as np
from ai_edge_litert.interpreter import Interpreter

# Mimi decode: 32 codes per frame -> latent -> 24 kHz audio, one shot over 256 frames.
dequant = Interpreter(model_path="dia2_mimi_dequant.tflite"); dequant.allocate_tensors()
decode = Interpreter(model_path="dia2_mimi_decode_t256.tflite"); decode.allocate_tensors()

latents = np.zeros((1, 512, 256), np.float32)          # tail stays zeroed: the path is causal
for tau in range(num_frames):                          # aligned (undelayed) codes, [32, num_frames]
    dequant.set_tensor(dequant.get_input_details()[0]["index"],
                       aligned[:, tau].reshape(1, 32, 1).astype(np.float32))
    dequant.invoke()
    latents[0, :, tau] = dequant.get_tensor(dequant.get_output_details()[0]["index"]).reshape(-1)

decode.set_tensor(decode.get_input_details()[0]["index"], latents)
decode.invoke()
audio = decode.get_tensor(decode.get_output_details()[0]["index"]).reshape(-1)[:num_frames * 1920]
```

## Performance and memory

On a Pixel 8a a 4-second utterance takes ~190 s: 71 warm-up frames (temporal only, ×2 guidance
branches) plus ~67 generated frames, each running 2 temporal steps and 2×31 depformer stages. The
process peaks at **~4.6 GB RSS** and settles around 3.2 GB — on an 8 GB phone, close other apps.

## Validation

Every ported component was checked against the reference implementation on host before it reached
the device:

| Component | Check | Result |
|---|---|---|
| StateMachine (multiplex + lookahead) | vs recorded reference frame stream | 0/60 mismatches |
| Tokenizer + `parse_script` | vs reference entries | 10/10 exact |
| Depformer 31-stage KV glue | vs torch depformer | corr 1.0000, argmax 31/31 |
| Mimi decode | vs torch full-sequence decode | corr 0.999999 |
| Voice-prefix warm-up | vs reference warm-up + generation | 0 mismatches (71 + 63 frames) |

## License

Apache-2.0, inherited from [nari-labs/Dia2-1B](https://huggingface.co/nari-labs/Dia2-1B). The Mimi
codec graphs derive from [kyutai/mimi](https://huggingface.co/kyutai/mimi) (CC-BY-4.0).