anira-project/nemotron-3-diarization

Nemotron 3 Diarization for anira

NVIDIA's Nemotron 3 Diarization, a Streaming Sortformer with a speaker cache that labels who spoke when (up to 8 speakers), for streaming inference with anira: a q4 ONNX export (mirrored unchanged), and ExecuTorch and Core ML exports of NVIDIA's checkpoint (export/). The speaker engine of anira foundation models (speech to text).

Files

File	Size	SHA256
onnx/model_q4.onnx	359,731	3e5e93dfd8289770343c402079d36034faa4a2eac260f66115f85ba9a2de41b2
onnx/model_q4.onnx_data	82,768,000	4513ed877d7cb83944bcb27c747221a7180e03f790d1c703e50af01792b7b8f4

ExecuTorch and Core ML

Exported from NVIDIA's checkpoint (transformers' Nemotron3DiarizationForAudioFrameClassification) with export/, as one streaming step with the tensor names of the ONNX export and fixed shapes. The state stays outside, as with the ONNX export: the caller keeps the speaker cache and FIFO and feeds them back.

Inputs	Outputs
input_features [1,104,128]: one chunk, 9 encoder frames scored and 4 look-ahead (low-latency mode), 8 mel frames each	logits [1,1640,8]: speaker logits of every frame
cached_embeds [1,192,512]: speaker cache and FIFO (128 + 64 frames), padded at the front while they fill	chunk_embeds [1,13,512]: the chunk's frames, for the FIFO
attention_mask [1,205]: 1 for every real frame, 0 for the padding	silence_embeds [512]

The mask is int64 (int32 in Core ML). Front padding keeps the real frames' rotary distances, and the padding is zeroed before the head's 3-frame upsampling convolution, so a padded step equals an unpadded one.

File	Engine	Size	SHA256
executorch/nemotron_diarization.pte	ExecuTorch 1.3.1, CPU (XNNPACK)	397,401,760	dc3cefafb4edcf435de64336090d3791599280e1c475e59d6e2b43e1b1cafdb7
executorch/nemotron_diarization_coreml.pte	ExecuTorch 1.3.1 (Core ML, Apple)	399,035,677	9cae034b6ea242e372a811b81e81f61308c19401d083da24fb6f8b0b4960ae8c
coreml/nemotron_diarization.mlpackage	Core ML (native)	398,485,283	per file in the catalog

Format	Precision	Per 90 ms chunk (M-series Mac)	Max speaker-probability deviation from transformers
ExecuTorch, CPU (XNNPACK)	float32	226 ms	3.3e-6
ExecuTorch, Core ML	float32	14 ms	3.5e-6
Core ML ML program (macOS 13, iOS 16)	float32	14 ms	3.9e-6
ONNX Runtime (this repository's q4)	4-bit weights	190 ms	3.0e-2

Measured over a 25 s recording of several voices, 34 chunks, each format with its own speaker cache fed back, against transformers' streaming forward with the same cache sizes. On a CPU, a 192-frame cache is slower than real time with every format here; Core ML runs on the GPU and Neural Engine. The CPU program also runs on the ExecuTorch runtime anira ships (anira-project/backends v2.4.0).

cd export
uv run export_executorch.py <folder> [--coreml] <audio.wav>...
uv run export_coreml.py <folder> <audio.wav>...
uv run check.py <audio.wav>...        # the ONNX export against transformers

Use with anira

Pin a revision and check each file's SHA256 (listed above) when you download:

https://huggingface.co/anira-project/nemotron-3-diarization/resolve/<revision>/onnx/model_q4.onnx

Licence

OpenMDW-1.1: LICENSE.OpenMDW-1.1, copied from OpenMDW/OpenMDW @ b26b32b. The model is NVIDIA's Nemotron 3 Diarization; the ONNX export is onnx-community/Nemotron-3-Diarization-ONNX @ 353b6f8, whose model card is kept here as UPSTREAM_MODEL_CARD.md with its notices of origin, as the licence asks. The ExecuTorch and Core ML files are exported from NVIDIA's checkpoint @ f667ed7 (export/, Apache-2.0).