mlboydaisuke/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

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)

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:

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. The Mimi codec graphs derive from kyutai/mimi (CC-BY-4.0).