litert-community/Speaker-Diarization-LiteRT

Speaker Diarization — LiteRT on-device stack (pyannote 3.1 recipe)

Who-spoke-when timeline from the on-device pipeline (2-speaker conversation). Colored bars = per-speaker turns.

On-device speaker diarization ("who spoke when") for Android, following the pyannote/speaker-diarization-3.1 recipe (MIT):

file	model	runtime	license
wespeaker_emb_fp16.tflite	WeSpeaker ResNet34 speaker embedding (6.6 M)	LiteRT CompiledModel GPU	CC-BY-4.0 (weights)
pyannote_seg30.onnx	pyannote segmentation-3.0 (PyanNet SincNet+BiLSTM, 1.5 M)	onnxruntime CPU	MIT

The segmentation BiLSTM has no mobile-GPU kernel, so it runs on CPU (tiny and fast); the heavy embedding CNN runs fully on the GPU. Verified on a Pixel 8a (Tensor G3): embedding 108 / 108 nodes LITERT_CL (full residency, 1 partition), ~1.2 ms per window, device-vs-PyTorch cosine 0.99997 (fp16, 13.4 MB); segmentation ONNX corr 1.0 / per-frame argmax agreement 100% vs PyTorch.

I/O

Embedding wespeaker_emb_fp16.tflite

• Input [1, 500, 80] float32 — kaldi log-mel fbank (25 ms / 10 ms hamming, 80 bins, dither 0, waveform ×2¹⁵ before fbank), CMN'd (subtract the per-bin mean over the 500 frames). 500 frames = 80 240 samples = 5.015 s @ 16 kHz; tile-pad shorter speech.
• Output [1, 256] — speaker embedding (L2-normalize before cosine comparison).

Segmentation pyannote_seg30.onnx

• Input [1, 1, 160000] float32 — 10 s @ 16 kHz mono, [-1, 1].
• Output [1, 589, 7] — per-frame log-probs over the powerset classes {∅, s1, s2, s3, s1s2, s1s3, s2s3} (≤3 local speakers, ≤2 concurrent).

Minimal usage (Python)

import numpy as np, soundfile as sf, torch, onnxruntime as ort
import torchaudio.compliance.kaldi as kaldi
from ai_edge_litert.interpreter import Interpreter

wav, sr = sf.read("speech.wav", dtype="float32")     # 16 kHz mono, [-1, 1]

# 1) segmentation: who is active in a 10 s window
seg = ort.InferenceSession("pyannote_seg30.onnx")
x = np.zeros(160000, np.float32); n = min(len(wav), 160000); x[:n] = wav[:n]
ps = seg.run(None, {"waveform": x[None, None]})[0][0]          # [589, 7] log-probs
PS = [(), (0,), (1,), (2,), (0, 1), (0, 2), (1, 2)]            # powerset classes
active = [PS[c] for c in ps.argmax(1)]                          # local speakers per ~17 ms frame

# 2) speaker embedding of a 5.015 s snippet (tile-pad shorter speech to 80240 samples)
snip = np.resize(wav, 80240).astype(np.float32)
fb = kaldi.fbank(torch.tensor(snip[None]) * 32768, num_mel_bins=80, frame_length=25.0,
                 frame_shift=10.0, dither=0.0, window_type="hamming", sample_frequency=16000)
fb = (fb - fb.mean(0, keepdim=True))[None].numpy()              # CMN -> [1, 500, 80]
emb = Interpreter(model_path="wespeaker_emb_fp16.tflite"); emb.allocate_tensors()
emb.set_tensor(emb.get_input_details()[0]["index"], fb); emb.invoke()
e = emb.get_tensor(emb.get_output_details()[0]["index"])[0]     # [256]
e /= np.linalg.norm(e)                                          # cosine-compare across snippets,
                                                                # cluster at distance 0.7046

Kotlin (Android)

// Embedding — LiteRT CompiledModel GPU:  implementation("com.google.ai.edge.litert:litert:2.1.5")
val emb = CompiledModel.create(File(ctx.filesDir, "wespeaker_emb_fp16.tflite").absolutePath,
    CompiledModel.Options(Accelerator.GPU), null)
val eIn = emb.createInputBuffers()
val eOut = emb.createOutputBuffers()
eIn[0].writeFloat(fbankCmn)            // [500 * 80]: kaldi fbank (25/10 ms hamming, x2^15) + CMN,
emb.run(eIn, eOut)                     //   see Fbank.kt in the speaker_diarization LiteRT sample
val e = eOut[0].readFloat()            // [256] — L2-normalize, cosine-compare, cluster at 0.7046

// Segmentation — onnxruntime CPU:  implementation("com.microsoft.onnxruntime:onnxruntime-android:1.24.3")
val env = OrtEnvironment.getEnvironment()
val seg = env.createSession(File(ctx.filesDir, "pyannote_seg30.onnx").absolutePath,
    OrtSession.SessionOptions())
OnnxTensor.createTensor(env, FloatBuffer.wrap(window), longArrayOf(1, 1, 160000)).use { t ->
    seg.run(mapOf(seg.inputNames.first() to t)).use { out ->
        @Suppress("UNCHECKED_CAST")
        val ps = (out[0].value as Array<Array<FloatArray>>)[0]   // [589][7] powerset log-probs
        // per-frame argmax -> {∅, s1, s2, s3, s1s2, s1s3, s2s3}
    }
}

Pipeline (as in the reference)

Sliding 10 s windows → powerset argmax → per-(window, local speaker) units with enough solo speech → embedding of each unit's concatenated solo audio → agglomerative clustering (centroid linkage, cosine distance, threshold 0.7046 from the 3.1 config) → stitched global timeline.

Conversion

Embedding converted with litert-torch from pyannote/wespeaker-voxceleb-resnet34-LM: a pure CNN (no maxpool stem) — zero re-authoring except the StatsPool standard deviation (down-scaled unbiased variance, fp16-safe). fp16 tflite vs PyTorch cosine 1.0000. Segmentation exported to ONNX from pyannote/segmentation-3.0.

Upstream

• pyannote.audio (MIT) — please cite Bredin 2023 (pyannote 2.x/3.x) when you use these models.
• WeSpeaker (Apache-2.0 code; the voxceleb-resnet34-LM weights are CC-BY-4.0).