feat(diarizer): Nemotron 3 Diarization support (8-speaker streaming Sortformer)

Documentation/ANE_Profiler.md

@@ -159,6 +159,51 @@ fused `decoder_joint` (B1).
 
 ---
 
+# Diarization (streaming)
+
+Compute plan (`Scripts/ane_profile.swift`, `--units all`) plus **warm per-call latency on real audio**
+(NVIDIA's 97.6 s 8-voice demo clip, M5 Pro, macOS 26.7; Nemotron via `nemotron3-diarize --profile`,
+Sortformer derived from wall RTFx over its 203 calls, so it includes host time and the cold first call).
+Both models are pure forward passes over `[speaker cache | FIFO | chunk]`; `T` is that packed length.
+
+| Model | Preset | Latency | Audio/call | T | ANE | GPU | CPU | ops | Size | ANE ms/call | GPU ms/call | ANE ms per audio-s |
+|-------|--------|--------:|-----------:|--:|----:|----:|----:|----:|-----:|------------:|------------:|-------------------:|
+| Sortformer v2.1 | fast | 1.04 s | 0.48 s | 242 | 94% | 0% | 6% | 1526 | 229 MB | 12.5 | 10.4 | 26 |
+| Sortformer v2.1 | high context | 30.4 s | 27.2 s | — | 94% | 0% | 6% | 1526 | 243 MB | — | — | — |
+| Sortformer v2.1 | offline (fused) | 30.7 s | 30.7 s | — | 99% | 0% | 1% | 1497 | 230 MB | — | — | — |
+| Nemotron 3 | low | 1.04 s | 0.72 s | 541 | 98% | 0% | 2% | 1178 | 190 MB | 27.1 | 12.0 | 38 |
+| Nemotron 3 | fast32 | 2.88 s | 2.56 s | 340 | 98% | 0% | 2% | 1178 | 190 MB | 11.6 | 12.6 | 4.5 |
+| Nemotron 3 | fast128 | 10.56 s | 10.24 s | 436 | 98% | 0% | 2% | 1178 | 190 MB | 20.6 | 37.0 | 2.0 |
+| Nemotron 3 | offline | 30.4 s | 27.2 s | 684 | 98%* | 0% | 2% | 1178 | 190 MB | fails* | 11.2 | — |
+| Nemotron 3 | fast32-split-w8a8 | 2.88 s | 2.56 s | 340 | 100% | 0% | 0% | 1649 | 95 MB | 9.7 | — | 3.8 |
+| Nemotron 3 | c128-split-w8a8 | 10.56 s | 10.24 s | 436 | 100% | 0% | 0% | 1649 | 95 MB | ~18 | — | 1.8 |
+
+\* `MLComputePlan` reports the placement CoreML *intends*; Nemotron 3 `offline` (3040 mel frames) fails
+`ANECCompile` at runtime and silently runs on the GPU. Chunk mel input ≤ 1376 frames compiles for the
+ANE; 1440+ does not. The split-graph presets bypass the cliff (host does feature stacking + the
+1024→512 projection).
+
+**Reading the table**
+
+- Sortformer's 2% CPU residue and Nemotron's 2% are index/gather ops around the state packing; the
+  split-graph variants move that packing to the host and leave a pure-fp transformer that is 100% ANE.
+- **Per-call cost scales with `T`, not with audio advanced.** Nemotron `low` (T=541) costs 2.2× Sortformer
+  fast (T=242) per call on the ANE and advances 1.5× the audio; at 1.04 s latency the two are within 1.5×
+  of each other per audio-second. Bigger Nemotron chunks amortize the fixed state: fast32 is 8× cheaper
+  than `low` per audio-second at higher DER-neutral latency, fast128 19× cheaper.
+- **The M5 Pro GPU beats the ANE at 1.04 s latency for both models** (Nemotron `low` 12.0 vs 27.1 ms,
+  Sortformer fast 10.4 vs 12.5) — ANE tiling of these packed sequences is unfavourable — while the ANE
+  wins for fast128 (20.6 vs 37.0). `.all` picks per-op, not
+  per-model, so choose the route explicitly for `low` on Macs; on iPhone the ANE is the only fast route.
+- Sortformer's first GPU run on a fresh process paid a ~2.4 s cold compile on call 1 (whole-clip RTFx
+  22× instead of 46×); the table's GPU figure is the warm second run. Its ANE cold cost is small.
+  Nemotron's cold ANE compile is ~1 s for the monolithic presets.
+- Sequence length is the cost: zero-shot layer drops, W8A8 on the monolithic graph, batch>1 on the ANE
+  and speaker-cache/FIFO shrinking were all measured and rejected (see the Nemotron 3 conversion notes);
+  the remaining lever is reusing the static state's attention across speaker-cache updates, untested.
+
+---
+
 # Diarization (offline)
 
 | Pipeline | Type | Chunk | ANE | GPU | CPU | ops | Size |

Documentation/Diarization/Nemotron3.md

@@ -0,0 +1,124 @@
+# Nemotron 3 Diarization
+
+FluidAudio support for NVIDIA's **Nemotron 3 Diarization** (streaming Sortformer
+successor): up to **8 speakers**, arrival-order speaker channels, 10 ms output
+resolution, streaming and offline profiles from a single checkpoint.
+
+> **Model availability:** the converted CoreML presets are published at
+> [`FluidInference/nemotron-3-diarization-coreml`](https://huggingface.co/FluidInference/nemotron-3-diarization-coreml)
+> (gated until NVIDIA's public release — request access, then set `HF_TOKEN`;
+> ungated afterwards). `Nemotron3Models.loadFromHuggingFace` downloads one preset
+> bundle on first use; `Nemotron3Models.load(config:directory:)` loads a local copy.
+
+## Quick start
+
+```swift
+import FluidAudio
+
+let config = Nemotron3Config.fast32  // recommended default
+let models = try await Nemotron3Models.loadFromHuggingFace(config: config)
+let diarizer = Nemotron3Diarizer(config: config, models: models)
+
+let (probs, frames) = try diarizer.processComplete(audioSamples)  // 16 kHz mono
+let segments = Nemotron3Diarizer.segments(probabilities: probs, frameCount: frames)
+// arrival-ordered speaker segments at 10 ms resolution, up to 8 speakers
+```
+
+Optional VAD gating for silence-heavy audio (skips inference over non-speech while
+preserving the output timeline):
+
+```swift
+let (probs, frames) = try diarizer.processComplete(audioSamples, speechMask: mask)
+```
+
+## Streaming (microphone / live audio)
+
+Feed 16 kHz samples as they arrive; a chunk runs as soon as its right context is
+buffered (`config.latencySeconds` after the chunk starts) and returns
+`config.chunkSeconds` of new 10 ms frames. The path is frame-exact with
+`processComplete` on the same audio.
+
+```swift
+diarizer.reset()
+for samples in microphoneSteps {            // any granularity, e.g. 320 ms
+    diarizer.appendAudio(samples)
+    for chunk in try diarizer.processBufferedAudio() {
+        timeline.append(contentsOf: chunk.probabilities)   // [frames * 8]
+    }
+}
+for chunk in try diarizer.finishStream() {  // flushes the trailing partial chunk
+    timeline.append(contentsOf: chunk.probabilities)
+}
+```
+
+`Nemotron3Diarizer` is not thread-safe: own it from one actor or task. For
+word→speaker attribution pair it with
+`StreamingUnifiedAsrManager.consumeWordTimings()` and pick, per word, the speaker
+slot with the most activity over the word's span.
+
+## Choosing a preset
+
+Latency = (chunk + right context) x 80 ms — the audio buffered before a result is
+final. Audio chunk = new audio consumed per model call; larger chunks amortize the
+fixed speaker-cache cost, which *improves* accuracy while increasing throughput.
+
+| Preset | Size | Audio chunk/call | Latency | Pros | Cons |
+|---|---|---|---|---|---|
+| `low` | 190 MB | 0.72 s | 1.04 s | Best quality at real streaming latency; NVIDIA's reference config | Heaviest ANE use per second of audio |
+| `fast` | 190 MB | 0.72 s | 1.04 s | ~3x cheaper per call than `low` — leaves ANE room for concurrent ASR | Slightly lower accuracy than `low` |
+| `fast32` | 190 MB | 2.56 s | 2.88 s | **Recommended default** — `low`-level accuracy at near-`fast` cost | Latency too high for live-caption UX |
+| `fast128` | 190 MB | 10.24 s | 10.56 s | Best accuracy of the streaming lineup; highest streaming throughput | Near-live only; results trail by ~10 s |
+| `offline` | 190 MB | 27.2 s | 30.4 s | Highest accuracy; fastest batch profile | GPU-only (ANE compiler limit); 30 s latency |
+| `s32-split-w8a8`* | **95 MB** | 2.56 s | 2.88 s | Half size, 100% ANE-resident graph, zero GPU use — the iOS pick | Requires `pre_encode_proj_t.bin` alongside the model |
+| `c128-split-w8a8`* | **95 MB** | 10.24 s | 10.56 s | Batch throughput without touching the GPU | Same split-mode requirement; ~10 s latency |
+
+\* Split-graph mode (`splitGraph` config flag): feature stacking and the 1024→512
+projection run host-side (one reshape + one `cblas_sgemm`), leaving a pure
+floating-point transformer graph that is fully ANE-resident and quantizes cleanly
+to W8A8. `Nemotron3Models.runSplit` handles the host-side work transparently.
+
+Quick chooser: hard ~1 s latency → `fast` (sharing the ANE with ASR) or `low`
+(diarizer owns the ANE) · general use → `fast32` · latency-flexible quality →
+`fast128` · recorded archives on a Mac → `offline` · iPhone/iPad, battery, or
+GPU-busy systems → the `split-w8a8` pair.
+
+Additional card profiles (`verylow`, `ultra`) and intermediate configurations exist
+via `Nemotron3Config.preset(named:)` / custom initializers but are dominated by the
+presets above for typical use.
+
+## CLI
+
+```bash
+# Diarize a file (prints segments; --output writes RTTM). Downloads the preset on first use.
+swift run fluidaudiocli nemotron3-diarize audio.wav --variant fast32
+
+# Same audio through the live path (appendAudio / processBufferedAudio / finishStream)
+swift run fluidaudiocli nemotron3-diarize audio.wav --variant fast32 --streaming
+
+# Benchmark against AMI / VoxConverse harnesses
+swift run fluidaudiocli nemotron3-benchmark --variant fast32 --collar 0
+
+# Batch processing with concurrent GPU workers
+swift run fluidaudiocli nemotron3-batch --workers 2 --files a,b,c
+```
+
+Useful flags: `--models <dir>` (local bundles instead of the HF download; required for
+`verylow`/`ultra` and sweep variants), `--compute-units ane|gpu|all`, `--profile`
+(per-stage wall breakdown), `--vad` (Silero-gated processing), sweep flags
+(`--chunk-len`, `--fifo`, `--spkcache`, `--rc`, `--update-period`) for
+custom-converted models.
+
+## Implementation notes
+
+- **State lives host-side**: the CoreML model is a pure forward pass over
+  `[speaker cache | FIFO | chunk]`; `Nemotron3StateUpdater` ports NeMo's
+  `streaming_update_async` (cache compression, learned silence embedding, FIFO
+  eviction) in Swift. Closed-loop output matches the NeMo reference at 99.995%
+  frame agreement on real audio.
+- Model outputs are fp16 with padded rows; readback uses a stride-aware
+  `vDSP_mmov` compaction (naive reads silently scramble or run ~40x slower —
+  see `Nemotron3TensorLayoutTests`).
+- Long ANE-route runs require the per-chunk autoreleasepool in `processComplete`
+  (IOSurface-backed outputs otherwise exhaust the pool after thousands of calls).
+- The mel frontend is the shared `AudioMelSpectrogram` (128 mel, 10 ms hop,
+  no normalization) — the same family as the Nemotron ASR models.

Scripts/materialize_alimeeting_card_audio.py

@@ -0,0 +1,83 @@
+#!/usr/bin/env python3
+"""Materialize AliMeeting Test audio for the NVIDIA Nemotron 3 card protocol.
+
+Card conditions (model card, Evaluation Datasets):
+  - AliMeeting Test Far  = far-field array audio  -> channel 0 of the 8-channel wav
+  - AliMeeting Test Near = "mix of headset microphones" -> equal-weight average of
+    the per-speaker N_SPK*.wav headset channels
+
+Inputs  : ~/FluidAudioDatasets/alimeeting/Test_Ali/Test_Ali_{far,near}/audio_dir
+Outputs : ~/FluidAudioDatasets/alimeeting/card/{far_ch0,near_mix}/<meeting>.wav
+          where <meeting> matches the nttcslab-sp/diar-forced-alignment RTTM names
+          (e.g. R8002_M8002).
+
+Idempotent: existing outputs are skipped.
+"""
+
+import re
+import subprocess
+import sys
+from collections import defaultdict
+from pathlib import Path
+
+ROOT = Path.home() / "FluidAudioDatasets" / "alimeeting"
+FAR_IN = ROOT / "Test_Ali" / "Test_Ali_far" / "audio_dir"
+NEAR_IN = ROOT / "Test_Ali" / "Test_Ali_near" / "audio_dir"
+FAR_OUT = ROOT / "card" / "far_ch0"
+NEAR_OUT = ROOT / "card" / "near_mix"
+
+MEETING_RE = re.compile(r"^(R\d+_M\d+)")
+
+
+def run(cmd):
+    subprocess.run(cmd, check=True, capture_output=True)
+
+
+def materialize_far():
+    FAR_OUT.mkdir(parents=True, exist_ok=True)
+    for wav in sorted(FAR_IN.glob("*.wav")):
+        m = MEETING_RE.match(wav.stem)
+        if not m:
+            print(f"skip (unrecognized name): {wav.name}")
+            continue
+        out = FAR_OUT / f"{m.group(1)}.wav"
+        if out.exists():
+            continue
+        # Channel 0 of the far-field array, 16 kHz mono.
+        run([
+            "ffmpeg", "-nostdin", "-v", "error", "-i", str(wav),
+            "-af", "pan=mono|c0=c0", "-ar", "16000", "-c:a", "pcm_s16le", str(out),
+        ])
+        print(f"far  {out.name}")
+
+
+def materialize_near():
+    NEAR_OUT.mkdir(parents=True, exist_ok=True)
+    groups = defaultdict(list)
+    for wav in sorted(NEAR_IN.glob("*.wav")):
+        m = MEETING_RE.match(wav.stem)
+        if m:
+            groups[m.group(1)].append(wav)
+    for meeting, wavs in sorted(groups.items()):
+        out = NEAR_OUT / f"{meeting}.wav"
+        if out.exists():
+            continue
+        # Equal-weight average of the headset channels: amix with default
+        # normalize=1 divides the sum by the input count.
+        cmd = ["ffmpeg", "-nostdin", "-v", "error"]
+        for wav in wavs:
+            cmd += ["-i", str(wav)]
+        cmd += [
+            "-filter_complex", f"amix=inputs={len(wavs)}:duration=longest",
+            "-ar", "16000", "-c:a", "pcm_s16le", str(out),
+        ]
+        run(cmd)
+        print(f"near {out.name} ({len(wavs)} headsets)")
+
+
+if __name__ == "__main__":
+    if not FAR_IN.is_dir() or not NEAR_IN.is_dir():
+        sys.exit(f"AliMeeting Test_Ali audio not found under {ROOT}")
+    materialize_far()
+    materialize_near()
+    print(f"done: {len(list(FAR_OUT.glob('*.wav')))} far, {len(list(NEAR_OUT.glob('*.wav')))} near")

Scripts/materialize_notsofar_card_audio.py

@@ -0,0 +1,120 @@
+#!/usr/bin/env python3
+"""Materialize NOTSOFAR1 eval audio + references for Nemotron 3 card-style rows.
+
+Card conditions (NVIDIA model card, Evaluation Datasets):
+  - NOTSOFAR1 Eval MHM = "mix of headset microphones" -> equal-weight average of
+    close_talk/CT_*.wav per meeting
+  - NOTSOFAR1 Eval SC  = "far-field single-channel"   -> ch0.wav of one
+    single-channel device per meeting (first sc_* directory sorted by name;
+    NVIDIA's exact device/session list is unpublished)
+
+References: NVIDIA scored against unpublished FastMSS forced alignments. We build
+RTTMs from the released gt_transcription.json word timings, merging consecutive
+same-speaker words when the inter-word gap is <= 0.2 s (mirrors the
+nttcslab-sp/diar-forced-alignment word-alignment convention used for AMI and
+AliMeeting). Our NOTSOFAR rows are therefore protocol-adjacent, not
+protocol-identical — same audio conditions and scoring settings, different
+reference timing source.
+
+Inputs  : ~/FluidAudioDatasets/notsofar/hf/benchmark-datasets/eval_set/240825.1_eval_full_with_GT/MTG/MTG_*
+Outputs : ~/FluidAudioDatasets/notsofar/card/{eval_mhm,eval_sc}/<meeting>.wav
+          ~/FluidAudioDatasets/notsofar/card/rttm/<meeting>.rttm
+"""
+
+import json
+import subprocess
+import sys
+from pathlib import Path
+
+ROOT = Path.home() / "FluidAudioDatasets" / "notsofar"
+MTG_ROOT = ROOT / "hf" / "benchmark-datasets" / "eval_set" / "240825.1_eval_full_with_GT" / "MTG"
+MHM_OUT = ROOT / "card" / "eval_mhm"
+SC_OUT = ROOT / "card" / "eval_sc"
+RTTM_OUT = ROOT / "card" / "rttm"
+
+WORD_MERGE_GAP = 0.2  # seconds
+
+
+def run(cmd):
+    subprocess.run(cmd, check=True, capture_output=True)
+
+
+def build_rttm(meeting_dir: Path, meeting: str) -> bool:
+    gt_path = meeting_dir / "gt_transcription.json"
+    if not gt_path.exists():
+        return False
+    utterances = json.loads(gt_path.read_text())
+
+    # Word-level segments per speaker, merged at <= WORD_MERGE_GAP gaps.
+    words = []
+    for utt in utterances:
+        spk = utt["speaker_id"]
+        timing = utt.get("word_timing") or []
+        for _, start, end in timing:
+            words.append((spk, float(start), float(end)))
+        if not timing:
+            words.append((spk, float(utt["start_time"]), float(utt["end_time"])))
+    words.sort(key=lambda w: (w[0], w[1]))
+
+    segments = []
+    for spk, start, end in words:
+        if segments and segments[-1][0] == spk and start - segments[-1][2] <= WORD_MERGE_GAP:
+            segments[-1][2] = max(segments[-1][2], end)
+        else:
+            segments.append([spk, start, end])
+    segments.sort(key=lambda s: s[1])
+
+    lines = [
+        f"SPEAKER {meeting} 1 {start:.3f} {end - start:.3f} <NA> <NA> {spk} <NA> <NA>"
+        for spk, start, end in segments
+        if end > start
+    ]
+    (RTTM_OUT / f"{meeting}.rttm").write_text("\n".join(lines) + "\n")
+    return True
+
+
+def materialize():
+    for d in (MHM_OUT, SC_OUT, RTTM_OUT):
+        d.mkdir(parents=True, exist_ok=True)
+
+    meetings = sorted(p for p in MTG_ROOT.glob("MTG_*") if p.is_dir())
+    if not meetings:
+        sys.exit(f"no meetings found under {MTG_ROOT}")
+
+    n_mhm = n_sc = 0
+    for meeting_dir in meetings:
+        meeting = meeting_dir.name
+        if not build_rttm(meeting_dir, meeting):
+            print(f"skip {meeting}: no gt_transcription.json")
+            continue
+
+        mhm_out = MHM_OUT / f"{meeting}.wav"
+        ct_wavs = sorted((meeting_dir / "close_talk").glob("CT_*.wav"))
+        if ct_wavs and not mhm_out.exists():
+            cmd = ["ffmpeg", "-nostdin", "-v", "error"]
+            for wav in ct_wavs:
+                cmd += ["-i", str(wav)]
+            cmd += [
+                "-filter_complex", f"amix=inputs={len(ct_wavs)}:duration=longest",
+                "-ar", "16000", "-c:a", "pcm_s16le", str(mhm_out),
+            ]
+            run(cmd)
+            n_mhm += 1
+
+        sc_out = SC_OUT / f"{meeting}.wav"
+        sc_devices = sorted(d for d in meeting_dir.glob("sc_*") if (d / "ch0.wav").exists())
+        if sc_devices and not sc_out.exists():
+            run([
+                "ffmpeg", "-nostdin", "-v", "error", "-i", str(sc_devices[0] / "ch0.wav"),
+                "-ar", "16000", "-c:a", "pcm_s16le", str(sc_out),
+            ])
+            n_sc += 1
+
+    print(
+        f"meetings {len(meetings)}: mhm {len(list(MHM_OUT.glob('*.wav')))} "
+        f"sc {len(list(SC_OUT.glob('*.wav')))} rttm {len(list(RTTM_OUT.glob('*.rttm')))}"
+    )
+
+
+if __name__ == "__main__":
+    materialize()