pyannote-speaker-diarization-onnx

pyannote/speaker-diarization-community-1 — ONNX

ONNX export of the two neural components of pyannote/speaker-diarization-community-1: the PyanNet segmentation model and the WeSpeakerResNet34 speaker-embedding model. Exported with torch.onnx.export (legacy TorchScript exporter, opset 17), validated numerically against the original PyTorch checkpoints.

Available Files

No FP16 for the embedding model — see Known Limitations.

Usage

Segmentation

import numpy as np
import onnxruntime as ort
import soundfile as sf

audio, sr = sf.read("audio.wav")
assert sr == 16000, "resample to 16kHz first"
waveform = audio.astype(np.float32)[None, None, :]  # (1, 1, samples)

sess = ort.InferenceSession("segmentation/model.onnx")
scores = sess.run(None, {"waveform": waveform})[0]  # (1, frames, 7)

scores are log-probabilities over the powerset encoding of {speaker#1, speaker#2, speaker#3} — 7 classes: silence, each speaker alone, and each pair overlapping. Frame rate is ~58.8 Hz (a 10s chunk produces 589 frames). Take argmax(-1) per frame to get the active class.

Speaker embedding

The embedding model's fbank frontend (torch.vmap over torchaudio.compliance.kaldi.fbank) isn't traceable to ONNX, so it's computed separately in Python before calling the exported ResNet:

import numpy as np
import onnxruntime as ort
import torch
import torchaudio.compliance.kaldi as kaldi

waveform = torch.from_numpy(audio.astype(np.float32))[None, None, :]  # (1, 1, samples)

fbank = kaldi.fbank(
    waveform.squeeze(0), num_mel_bins=80, frame_length=25, frame_shift=10,
    dither=0.0, sample_frequency=16000, window_type="hamming", use_energy=False,
)
fbank = (fbank - fbank.mean(dim=0, keepdim=True))[None]  # (1, frames, 80)

sess = ort.InferenceSession("embedding/model.onnx")
embedding = sess.run(None, {"fbank": fbank.numpy()})[0]  # (1, 256)

Validation

Checked against pyannote.audio.core.model.Model.from_pretrained(...) (the original PyTorch checkpoints):

Segmentation is consumed via argmax (powerset → binary speaker activity), so the raw-logit deltas in fp16/int8 don't change predictions on the inputs tested. Embeddings are consumed via cosine similarity for clustering, so 0.99 similarity — not raw magnitude — is the number that matters; validate against your own audio before relying on int8 in production.

Known Limitations

• No FP16 embedding model. onnxconverter-common's convert_float_to_float16 leaves a real fp16/fp32 type mismatch in the stats-pooling subgraph's unbiased-variance correction (a Sub/Cast pair computing N-1 frames), which onnxruntime then refuses to load. Blocking the offending nodes (and the whole stats_pool subgraph) via node_block_list just moves the mismatch to a different node at the block boundary. FP32 and INT8 are provided instead.
• Clustering is not included. VBx clustering (PLDA + variational Bayes HMM resegmentation) runs outside the neural network, on top of the embeddings — reimplement with numpy or use pyannote.audio.
• Legacy TorchScript exporter required. torch.onnx.export's dynamo/torch.export default (torch ≥ 2.9) fails on PyanNet's LSTM. Exported with dynamo=False.

License

CC-BY-4.0, inherited from the base model.