| --- |
| license: cc-by-4.0 |
| library_name: coreml |
| base_model: |
| - nvidia/diar_streaming_sortformer_4spk-v2.1 |
| base_model_relation: finetune |
| tags: |
| - speaker-diarization |
| - speech |
| - audio |
| - coreml |
| - apple |
| - ios |
| - macos |
| - sortformer |
| - streaming |
| pipeline_tag: automatic-speech-recognition |
| --- |
| # Sortformer CoreML Models |
|
|
| Streaming speaker diarization models converted from NVIDIA's Sortformer to CoreML for Apple Silicon. |
|
|
| ## Model Variants |
|
|
| | Variant | File | Latency | Use Case | |
| | -------------------- | ------------------------------------ | ------- | --------------------- | |
| | **Fastest v2** | `Sortformer_v2.mlmodelc` | ~1.04s | Low latency streaming | |
| | **Fastest v2.1** | `Sortformer_v2.1.mlmodelc` | ~1.04s | Low latency streaming | |
| | **NVIDIA Low v2** | `SortformerNvidiaLow_v2.mlmodelc` | ~1.04s | Low latency streaming | |
| | **NVIDIA Low v2.1** | `SortformerNvidiaLow_v2.1.mlmodelc` | ~1.04s | Low latency streaming | |
| | **NVIDIA High v2** | `SortformerNvidiaHigh_v2.mlmodelc` | ~30.4s | Best quality, offline | |
| | **NVIDIA High v2.1** | `SortformerNvidiaHigh_v2.1.mlmodelc` | ~30.4s | Best quality, offline | |
|
|
| The `v2` and `v2.1` refer to the version of the model weights to use. According to NVIDIA, `v2.1` is more robust in meeting scenarios. |
|
|
| ## Configuration Parameters |
|
|
| | Parameter | Default | NVIDIA Low | NVIDIA High | |
| | ------------------- | ------- | ---------- | ----------- | |
| | chunk_len | 6 | 6 | 340 | |
| | chunk_right_context | 7 | 7 | 40 | |
| | chunk_left_context | 1 | 1 | 1 | |
| | fifo_len | 40 | 188 | 40 | |
| | spkcache_len | 188 | 188 | 188 | |
| |
| ## Model Input/Output Shapes |
| |
| **General**: |
| |
| | Input | Shape | Description | |
| | ---------------- | --------------------- | ------------------------ | |
| | chunk | `[1, 8*(C+L+R), 128]` | Mel spectrogram features | |
| | chunk_lengths | `[1]` | Actual chunk length | |
| | spkcache | `[1, S, 512]` | Speaker cache embeddings | |
| | spkcache_lengths | `[1]` | Actual cache length | |
| | fifo | `[1, F, 512]` | FIFO queue embeddings | |
| | fifo_lengths | `[1]` | Actual FIFO length | |
|
|
| | Output | Shape | Description | |
| | ------------------------- | ------------------ | ------------------------------------- | |
| | speaker_preds | `[C+L+R+S+F, 4]` | Speaker probabilities (4 speakers) | |
| | chunk_pre_encoder_embs | `[C+L+R, 512]` | Embeddings for state update | |
| | chunk_pre_encoder_lengths | `[1]` | Actual embedding count | |
| | nest_encoder_embs | `[C+L+R+S+F, 192]` | Embeddings for speaker discrimination | |
| | nest_encoder_lengths | `[1]` | Actual speaker embedding count | |
| |
| Note: `C = chunk_len`, `L = chunk_left_context`, `R = chunk_right_context`, `S = spkcache_len`, `F = fifo_len`. |
|
|
| **Configuration-Specific Shapes**: |
| |
| | Input | Default | NVIDIA Low | NVIDIA High | |
| | ---------------- | --------------- | --------------- | ---------------- | |
| | chunk | `[1, 112, 128]` | `[1, 112, 128]` | `[1, 3048, 128]` | |
| | chunk_lengths | `[1]` | `[1]` | `[1]` | |
| | spkcache | `[1, 188, 512]` | `[1, 188, 512]` | `[1, 188, 512]` | |
| | spkcache_lengths | `[1]` | `[1]` | `[1]` | |
| | fifo | `[1, 40, 512]` | `[1, 188, 512]` | `[1, 40, 512]` | |
| | fifo_lengths | `[1]` | `[1]` | `[1]` | |
| |
| | Output | Default | NVIDIA Low | NVIDIA High | |
| | ------------------------- | --------------- | --------------- | --------------- | |
| | speaker_preds | `[1, 242, 128]` | `[1, 390, 128]` | `[1, 609, 128]` | |
| | chunk_pre_encoder_embs | `[1, 14, 512]` | `[1, 14, 512]` | `[1, 381, 512]` | |
| | chunk_pre_encoder_lengths | `[1]` | `[1]` | `[1]` | |
| | nest_encoder_embs | `[1, 242, 192]` | `[1, 390, 192]` | `[1, 609, 192]` | |
| | nest_encoder_lengths | `[1]` | `[1]` | `[1]` | |
|
|
| |
| | Metric | Default | NVIDIA High | |
| | ------------- | ------- | ----------- | |
| | Latency | ~1.12s | ~30.4s | |
| | RTFx (M4 Max) | ~5.7x | ~125.3x | |
|
|
| ## Usage with FluidAudio (Swift) |
|
|
| ```swift |
| import FluidAudio |
| |
| // Initialize with default config (auto-downloads from HuggingFace) |
| let diarizer = SortformerDiarizer(config: .default) |
| let models = try await SortformerModels.loadFromHuggingFace(config: .default) |
| diarizer.initialize(models: models) |
| |
| // Streaming processing |
| for audioChunk in audioStream { |
| if let result = try diarizer.processSamples(audioChunk) { |
| for frame in 0..<result.frameCount { |
| for speaker in 0..<4 { |
| let prob = result.getSpeakerPrediction(speaker: speaker, frame: frame) |
| } |
| } |
| } |
| } |
| |
| // Or batch processing |
| let timeline = try diarizer.processComplete(audioSamples) |
| for (speakerIndex, segments) in timeline.segments.enumerated() { |
| for segment in segments { |
| print("Speaker \(speakerIndex): \(segment.startTime)s - \(segment.endTime)s") |
| } |
| } |
| ``` |
| Performance |
|
|
| https://github.com/FluidInference/FluidAudio/blob/main/Documentation/Benchmarks.md |
|
|
| Files |
|
|
| Models |
|
|
| - Sortformer.mlpackage / .mlmodelc - Default config (low latency) |
| - SortformerNvidiaLow.mlpackage / .mlmodelc - NVIDIA low latency config |
| - SortformerNvidiaHigh.mlpackage / .mlmodelc - NVIDIA high latency config |
|
|
| Scripts |
|
|
| - convert_to_coreml.py - PyTorch to CoreML conversion |
| - streaming_inference.py - Python streaming inference example |
| - mic_inference.py - Real-time microphone demo |
|
|
| Source |
|
|
| Original model: https://huggingface.co/nvidia/diar_streaming_sortformer_4spk-v2.1 |
| |
| Credits & Acknowledgements |
| |
| This project would not have been possible without the significant technical contributions of https://huggingface.co/GradientDescent2718. |
| |
| Their work was instrumental in: |
| |
| - Architecture Conversion: Developing the complex PyTorch-to-CoreML conversion pipeline for the 17-layer Fast-Conformer and 18-layer Transformer heads. |
| - Build & Optimization: Engineering the static shape configurations that allow the model to achieve ~120x RTF on Apple Silicon. |
| - Logic Implementation: Porting the critical streaming state logic (speaker cache and FIFO management) to ensure consistent speaker identity tracking. |
| |
| This project was built upon the foundational work of the NVIDIA NeMo team. |
