mattt/diar-streaming-sortformer-4spk-v2-coreml

Streaming Sortformer 4spk v2 for Core ML

A Core ML conversion of NVIDIA's nvidia/diar_streaming_sortformer_4spk-v2 streaming speaker diarization model, for AnySpeechRecognizer. It predicts, for each 80 ms frame, an independent activity probability for each of up to four speakers.

Attribution and license

Streaming Sortformer 4spk v2 by NVIDIA (https://huggingface.co/nvidia/diar_streaming_sortformer_4spk-v2), licensed under CC BY 4.0 (https://creativecommons.org/licenses/by/4.0/). Converted to Core ML for AnySpeechRecognizer: the pre-encoder, FastConformer encoder, and Transformer are exported for one streaming step with a fixed speaker cache and FIFO, and the encoder and Transformer weights are stored in float16.

This repository is licensed under CC BY 4.0, the license of the base model. Redistributions must keep this attribution. This conversion is not the v2.1 checkpoint, which has a different license.

Files

File	Contents
SortformerLowLatency.mlmodelc	Compiled Core ML pipeline: pre-encoder (float32) and FastConformer encoder, Transformer, and speaker sigmoids (float16).
conversion.json	Source revision and digest, tool versions, settings, and the size and SHA-256 digest of every file.

Usage

The model runs one streaming step. The caller computes NeMo-compatible 128-bin log-mel features (16 kHz mono, 25 ms window, 10 ms hop, no normalization), keeps the speaker cache and FIFO queue between steps, and updates them with NeMo's SortformerModules.streaming_update from speaker_preds and chunk_pre_encoder_embs_out. AnySpeechRecognizer's SortformerDiarizationModel does all of this:

let model = SortformerDiarizationModel(assets: .compiledModel(url))
let diarizer = SpeakerDiarizer(model: model)

Pass the URL of the SortformerLowLatency.mlmodelc directory. Applications that use the model must show the attribution above.

Interface

Name	Type	Shape	Contents
chunk	float32	[1, 112, 128]	Log-mel features of the chunk with its left and right context, time-major, 128 bins
chunk_lengths	int32	[1]	Number of valid mel frames in chunk
spkcache	float32	[1, 188, 512]	Arrival-order speaker cache: pre-encoder embeddings of earlier frames
spkcache_lengths	int32	[1]	Number of valid frames in spkcache
fifo	float32	[1, 188, 512]	FIFO queue: pre-encoder embeddings of the most recent frames
fifo_lengths	int32	[1]	Number of valid frames in fifo
speaker_preds	float32	[1, 390, 4]	Speaker activity probabilities for the valid speaker cache, FIFO, and chunk frames, in that order
chunk_pre_encoder_embs_out	float16	[1, 14, 512]	Pre-encoder embeddings of the chunk, for the speaker cache and FIFO update
chunk_pre_encoder_lengths_out	int32	[1]	Number of valid frames in chunk_pre_encoder_embs_out

Each step takes 112 mel frames: 1 frame of left context, 6 chunk frames, and 7 frames of right context, each 8 mel frames (80 ms). Unused rows of chunk, spkcache, and fifo are ignored after their lengths. speaker_preds has the valid rows of the speaker cache, then of the FIFO, then of the chunk, followed by zeros.

Streaming configuration

NVIDIA's low-latency setting, with 1.04 s of algorithmic latency (chunk and right context). The values are also stored in the model's creator-defined metadata:

Key	Value
chunk_len	6
chunk_left_context	1
chunk_right_context	7
fifo_len	188
spkcache_len	188
spkcache_update_period	144
subsampling_factor	8
frame_duration	0.08
mel_feature_frames	48

Conversion

• Source: nvidia/diar_streaming_sortformer_4spk-v2 at revision 84edd514b8ef68004c10086918cd62f2148cbd59, diar_streaming_sortformer_4spk-v2.nemo (SHA-256 b371afce2c4958186469df33d939936b9746c89f38b10a69cfd2c61254e83329).
• The pre-encoder and the packing of cache, FIFO, and chunk frames are an ML program with float32 compute precision. The FastConformer encoder, Transformer, and speaker sigmoids are an ML program with float16 compute precision. Both use an iOS 16 and macOS 13 deployment target and form one Core ML pipeline.
• Tools: Python 3.12.12, PyTorch 2.7.0, NeMo 2.7.3, coremltools 9.0, macOS 27.0.1 (26A434).
• The conversion follows FluidInference/mobius (Apache-2.0) at commit 26f5fa8059d5ab3d34449adce6ff6ed86734e522.