Modelpublic

aufklarer/Sortformer-Diarization-CoreML

sourceHugging Facecc-by-4.0updated 2mo agoView on Hugging Face

1likes698downloads

Model Card

Sortformer Diarization (CoreML)

CoreML port of NVIDIA Sortformer for end-to-end speaker diarization on Apple Silicon. Runs on the Neural Engine. No separate embedding extraction or clustering — the model directly predicts per-frame speaker activity for up to 4 speakers, with streaming state (speaker cache + FIFO) carried across calls.

Model


Architecture	Sortformer (Sort Loss + 17-layer FastConformer + 18-layer Transformer)
Base model	`nvidia/diar_streaming_sortformer_4spk-v2.1` (117M params)
Precision	FP16 compute, FP32 I/O boundaries
Sample rate	16 kHz, 128 mel bins (n_fft 400, hop 160)
Max speakers	4 (arrival-order slots)
Speaker cache / FIFO	188 / 40 encoder frames
Frame duration	80 ms per prediction frame

Three variants ship from the same checkpoint, differing only in chunk shape:

Variant	New audio per call	Mel input	Use case
`Sortformer` (default)	~27 s	`[1, 3048, 128]`	offline batch, highest throughput
`Sortformer_balanced`	~8 s	`[1, 968, 128]`	~3× faster first-load compile
`Sortformer_streaming`	480 ms	`[1, 112, 128]`	incremental / realtime

Files

File	Size	Description
`Sortformer.mlmodelc` / `.mlpackage`	~230 MB	default variant (compiled / source)
`Sortformer_balanced.mlmodelc` / `.mlpackage`	~230 MB	balanced variant
`Sortformer_streaming.mlmodelc` / `.mlpackage`	~230 MB	streaming variant
`config.json`, `config_balanced.json`, `config_streaming.json`	<1 KB	chunk-shape hyperparameters per variant

Performance

Measured on M-series Apple Silicon (Neural Engine, warm):

Variant	Per-call latency	Realtime factor
default	one call per ~27 s of audio	~125–750×
balanced	one call per ~8 s	hundreds of ×
streaming	8.1 ms median per 480 ms step	~59× per step; ~39× end-to-end incremental

The streaming export is parity-gated: driven by NeMo's own streaming feature loader and cache-update algorithm, it matches NeMo's native forward_streaming loop at 100% frame-decision agreement (MAE 0.0005). On a five-file VoxConverse-dev pilot, the incremental session reaches 8.1% DER (collar 0.25) with correct speaker counts on all files.

Streaming interface

The streaming variant carries state through the CoreML interface; the speaker-cache update runs host-side between calls.

Inputs: chunk [1,112,128], chunk_lengths [1], spkcache [1,188,512], spkcache_lengths [1], fifo [1,40,512], fifo_lengths [1]

Outputs: speaker_preds_out [1,242,4], chunk_pre_encoder_embs_out [1,14,512], chunk_pre_encoder_lengths_out [1]

The pipeline contains two sub-models: PreEncoder (mel pre-encode + state concat) and Head (FastConformer + Transformer + sigmoid heads).

Usage

swift

// Whole-buffer diarization
let diarizer = try await SortformerDiarizer.fromPretrained()
let result = diarizer.diarize(audio: samples, sampleRate: 16000)

// Incremental streaming: push PCM in any size, stable speaker slots
let session = try await SortformerStreamingSession.fromPretrained()
let snapshot = try session.push(audio: samples)
let final = try session.finish()

bash

speech diarize meeting.wav --engine sortformer

Source

—Upstream model: nvidia/diar_streaming_sortformer_4spk-v2.1 · Sortformer paper · Streaming Sortformer paper
—Upstream license: NVIDIA Open Model License (weights); this conversion is published under CC-BY-4.0 with attribution to NVIDIA.