Dia2-1B-LiteRT

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

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 (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

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)

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)

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:

License

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

<!-- funnel:v1 -->

---

Want a different model on-device? Open a request — free, open weights only; the export and its measured numbers get published publicly.

<!-- /funnel:v1 -->