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

Documentation/Diarization/Nemotron3.md

@@ -0,0 +1,97 @@
+# 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 checkpoint is currently an early-access preview under
+> an NVIDIA evaluation license, so converted CoreML models are **not distributed**
+> with FluidAudio yet — they load from a local directory. HuggingFace auto-download
+> and full benchmark tables (DER / RTFx) will be published when NVIDIA's public
+> release lands.
+
+## Quick start
+
+```swift
+import FluidAudio
+
+let config = Nemotron3Config.fast32  // recommended default
+let models = try await Nemotron3Models.load(
+    config: config,
+    directory: localModelsDirectoryURL
+)
+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)
+```
+
+## 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)
+swift run fluidaudiocli nemotron3-diarize audio.wav --models <dir> --variant fast32
+
+# Benchmark against AMI / VoxConverse harnesses
+swift run fluidaudiocli nemotron3-benchmark --models <dir> --variant fast32 --collar 0
+
+# Batch processing with concurrent GPU workers
+swift run fluidaudiocli nemotron3-batch --models <dir> --workers 2 --files a,b,c
+```
+
+Useful flags: `--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()