diff --git a/Documentation/ANE_Profiler.md b/Documentation/ANE_Profiler.md index 9f536e981..65c752919 100644 --- a/Documentation/ANE_Profiler.md +++ b/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 | diff --git a/Documentation/Diarization/Nemotron3.md b/Documentation/Diarization/Nemotron3.md new file mode 100644 index 000000000..928dc0c3c --- /dev/null +++ b/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 ` (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. diff --git a/Scripts/materialize_alimeeting_card_audio.py b/Scripts/materialize_alimeeting_card_audio.py new file mode 100644 index 000000000..f91daed24 --- /dev/null +++ b/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}/.wav + where 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") diff --git a/Scripts/materialize_notsofar_card_audio.py b/Scripts/materialize_notsofar_card_audio.py new file mode 100644 index 000000000..4e7895101 --- /dev/null +++ b/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}/.wav + ~/FluidAudioDatasets/notsofar/card/rttm/.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} {spk} " + 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() diff --git a/Sources/FluidAudio/ASR/SenseVoice/SenseVoiceManager.swift b/Sources/FluidAudio/ASR/SenseVoice/SenseVoiceManager.swift index 6fcfb47b5..4203e8d2c 100644 --- a/Sources/FluidAudio/ASR/SenseVoice/SenseVoiceManager.swift +++ b/Sources/FluidAudio/ASR/SenseVoice/SenseVoiceManager.swift @@ -44,9 +44,28 @@ public actor SenseVoiceManager { /// Transcribe 16 kHz mono float samples (in [-1, 1]). public func transcribe(audio: [Float]) throws -> String { + try transcribeDetailed(audio: audio).text + } + + /// Transcribe and keep the model's leading query tags: detected language + /// (`zh`, `en`, `yue`, `ja`, `ko`, `nospeech`, …), emotion, and audio event. + public func transcribeDetailed(audio: [Float]) throws -> SenseVoiceTranscription { let features = try runPreprocessor(audio: audio) let (logits, validFrames) = try runEncoder(features: features) - return decode(logits: logits, validFrames: validFrames) + let raw = decodeRaw(logits: logits, validFrames: validFrames) + var tags: [String] = [] + let pattern = try NSRegularExpression(pattern: "<\\|([^|]*)\\|>") + let ns = raw as NSString + for m in pattern.matches(in: raw, range: NSRange(location: 0, length: ns.length)) { + tags.append(ns.substring(with: m.range(at: 1))) + } + let text = + raw + .replacingOccurrences(of: "<\\|[^|]*\\|>", with: "", options: .regularExpression) + .trimmingCharacters(in: .whitespaces) + let knownLanguages: Set = ["zh", "en", "yue", "ja", "ko", "nospeech"] + let language = tags.first { knownLanguages.contains($0) } ?? tags.first + return SenseVoiceTranscription(text: text, language: language, tags: tags) } // MARK: - Pipeline @@ -112,6 +131,13 @@ public actor SenseVoiceManager { /// Greedy CTC over the first `validFrames` (drop blank 0, collapse repeats), /// detokenize, then strip the `<|...|>` meta tags. private func decode(logits: MLMultiArray, validFrames: Int) -> String { + decodeRaw(logits: logits, validFrames: validFrames) + .replacingOccurrences(of: "<\\|[^|]*\\|>", with: "", options: .regularExpression) + .trimmingCharacters(in: .whitespaces) + } + + /// CTC decode with the `<|...|>` query tags left in place. + private func decodeRaw(logits: MLMultiArray, validFrames: Int) -> String { let frames = min(validFrames, logits.shape[1].intValue) // Per-frame argmax via the shared vDSP helper (~0.5s -> sub-ms for the // frames×vocab ~6.4M element scan), then CTC collapse (drop blank 0, @@ -124,10 +150,17 @@ public actor SenseVoiceManager { prev = best } - let raw = decodeCtcTokenIds(ids, vocabulary: models.vocabulary) - return - raw - .replacingOccurrences(of: "<\\|[^|]*\\|>", with: "", options: .regularExpression) - .trimmingCharacters(in: .whitespaces) + return decodeCtcTokenIds(ids, vocabulary: models.vocabulary) } } + +/// SenseVoice output with the model's leading query tags preserved. +public struct SenseVoiceTranscription: Sendable { + public let text: String + /// Detected language tag as emitted by the model (`zh`, `en`, `yue`, `ja`, `ko`, or + /// `nospeech`), `nil` when the model emitted none. Other languages are recognized + /// but tagged with the closest of the five. + public let language: String? + /// All leading tags in order: language, emotion, audio event, text-norm. + public let tags: [String] +} diff --git a/Sources/FluidAudio/Diarizer/Nemotron3/Nemotron3Diarizer.swift b/Sources/FluidAudio/Diarizer/Nemotron3/Nemotron3Diarizer.swift new file mode 100644 index 000000000..ce1498a11 --- /dev/null +++ b/Sources/FluidAudio/Diarizer/Nemotron3/Nemotron3Diarizer.swift @@ -0,0 +1,460 @@ +import Accelerate +import Foundation + +/// Streaming 8-speaker diarizer backed by NVIDIA's Nemotron 3 Diarization preview. +/// +/// Processes audio in fixed 80 ms-frame chunks through the CoreML forward pass and applies +/// NeMo's async speaker-cache/FIFO update host-side. Output is per-frame speaker activity +/// probability at 10 ms resolution, speaker slots ordered by first arrival. +/// +/// - Important: This class is **not** thread-safe. +public final class Nemotron3Diarizer { + + public let config: Nemotron3Config + private let models: Nemotron3Models + private let updater: Nemotron3StateUpdater + private var state: Nemotron3StreamingState + private let logger = AppLogger(category: "Nemotron3Diarizer") + + // Streaming (audio-in) state — see `appendAudio`. + private var frontend: Nemotron3StreamingFrontend + private var streamFinished = false + /// 10 ms output frames emitted so far by the streaming path. + public private(set) var streamedFrameCount = 0 + + /// Wall-time breakdown of the last `processComplete` call, in seconds. + public struct PipelineProfile: Sendable { + public var melSeconds: Double = 0 + public var chunkSliceSeconds: Double = 0 + public var inferenceSeconds: Double = 0 + public var inputPrepSeconds: Double = 0 + public var predictSeconds: Double = 0 + public var readbackSeconds: Double = 0 + public var stateUpdateSeconds: Double = 0 + public var outputAppendSeconds: Double = 0 + public var totalSeconds: Double = 0 + public var chunkCount: Int = 0 + /// Chunks skipped by VAD gating (no speech in the chunk's core window). + public var skippedChunks: Int = 0 + } + + /// Populated by `processComplete`; read after the call for stage-level analysis. + public private(set) var lastProfile = PipelineProfile() + + public init(config: Nemotron3Config, models: Nemotron3Models) { + self.config = config + self.models = models + self.updater = Nemotron3StateUpdater(config: config, silenceEmbedding: models.silenceEmbedding) + self.state = Nemotron3StreamingState(config: config) + self.frontend = Nemotron3StreamingFrontend(config: config) + } + + public func reset() { + state = Nemotron3StreamingState(config: config) + frontend.reset() + streamFinished = false + streamedFrameCount = 0 + } + + /// Process a complete audio buffer (16 kHz mono) and return per-frame speaker + /// probabilities at 10 ms resolution, [frames * 8] flattened. + /// Process a complete audio buffer. + /// + /// - Parameters: + /// - audio: 16 kHz mono samples. + /// - speechMask: Optional per-10 ms-frame speech mask (e.g. from `VadManager`). + /// Chunks whose core window contains no `true` frame skip inference entirely and + /// emit zero probabilities; streaming state does not advance across them (the + /// skipped region behaves like a pause in the stream). Callers should pre-pad + /// speech regions (~1 s) to protect onsets/offsets. + public func processComplete( + _ audio: [Float], speechMask: [Bool]? = nil + ) throws -> (probabilities: [Float], frameCount: Int) { + reset() + var profile = PipelineProfile() + let t0 = Date() + + var tStage = Date() + let mel = AudioMelSpectrogram() + let (featSeq, featLength, featSeqLength) = mel.computeFlatTransposed(audio: audio) + profile.melSeconds = Date().timeIntervalSince(tStage) + + var total = [Float]() + total.reserveCapacity(featLength * config.numSpeakers) + + var loader = Nemotron3FeatureLoader( + config: config, featSeq: featSeq, featLength: featLength, featSeqLength: featSeqLength) + let sub = config.subsamplingFactor + var coreStart = 0 + // Each chunk's prediction allocates IOSurface-backed output arrays; without a + // per-iteration autorelease drain, long ANE-route runs exhaust the IOSurface + // pool after a few thousand calls (issue #752 failure class). + while try autoreleasepool(invoking: { () -> Bool in + tStage = Date() + guard let chunk = loader.next() else { return false } + profile.chunkSliceSeconds += Date().timeIntervalSince(tStage) + + // VAD gate: emit zeros for speech-free chunks without running the model or + // advancing state. Output frame count must match the normal path exactly. + let coreEnd = min(coreStart + config.chunkLen * sub, featLength) + if let speechMask { + let lo = min(coreStart, speechMask.count) + let hi = min(coreEnd, speechMask.count) + let hasSpeech = lo < hi && speechMask[lo.. [Nemotron3ChunkResult] { + try runChunks(final: false) + } + + /// Flush the tail: pads the stream like `processComplete` does, runs the remaining + /// (possibly short) chunks, and trims the output to the audio's exact frame count. + /// The diarizer keeps its speaker state afterwards; call `reset()` before a new stream. + public func finishStream() throws -> [Nemotron3ChunkResult] { + guard !streamFinished else { return [] } + streamFinished = true + return try runChunks(final: true) + } + + private func runChunks(final: Bool) throws -> [Nemotron3ChunkResult] { + var results: [Nemotron3ChunkResult] = [] + while let chunk = frontend.nextChunk(final: final) { + var result = try autoreleasepool { + try step( + chunkFeatures: chunk.features, chunkMelLength: chunk.length, + leftOffsetMel: chunk.leftOffset, rightOffsetMel: chunk.rightOffset) + } + // Trim the tail chunk to the audio's exact 10 ms frame count (mirrors + // `processComplete`'s ceil(mel_frames) trim). + let remaining = frontend.melFramesComputed - streamedFrameCount + if final, result.frameCount > remaining { + let keep = max(remaining, 0) + result = Nemotron3ChunkResult( + probabilities: Array(result.probabilities[0..<(keep * config.numSpeakers)]), + frameCount: keep, numSpeakers: config.numSpeakers) + } + streamedFrameCount += result.frameCount + results.append(result) + } + return results + } + + /// Run one streaming step from raw mel features. + /// + /// - Parameters: + /// - chunkFeatures: Mel features [frames * 128] for lc+core+rc mel frames. + /// - chunkMelLength: Valid mel frames in `chunkFeatures`. + /// - leftOffsetMel: Mel frames of left context included at the start. + /// - rightOffsetMel: Mel frames of right context included at the end. + public func step( + chunkFeatures: [Float], + chunkMelLength: Int, + leftOffsetMel: Int, + rightOffsetMel: Int + ) throws -> Nemotron3ChunkResult { + let out = + config.splitGraph + ? try models.runSplit( + chunk: chunkFeatures, chunkLength: chunkMelLength, state: state, config: config) + : try models.run( + chunk: chunkFeatures, chunkLength: chunkMelLength, state: state, config: config) + let sub = config.subsamplingFactor + let lcEnc = (leftOffsetMel + sub / 2) / sub // round() + let rcEnc = (rightOffsetMel + sub - 1) / sub // ceil() + return try updater.update( + state: &state, + chunkEmbeddings: out.chunkEmbeddings, + chunkEncLength: out.chunkLength, + predictions: out.predictions, + highResPredictions: out.highResPredictions, + lc: lcEnc, + rc: rcEnc + ) + } + + /// Convert frame probabilities into arrival-ordered speaker segments. + public static func segments( + probabilities: [Float], frameCount: Int, numSpeakers: Int = 8, + threshold: Float = 0.5, frameSeconds: Float = 0.01, minDurationSeconds: Float = 0.2 + ) -> [Nemotron3Segment] { + var result: [Nemotron3Segment] = [] + for spk in 0.. threshold + if active, start == nil { + start = frame + } else if !active, let s0 = start { + let dur = Float(frame - s0) * frameSeconds + if dur >= minDurationSeconds { + result.append( + Nemotron3Segment( + speakerIndex: spk, + startSeconds: Float(s0) * frameSeconds, + endSeconds: Float(frame) * frameSeconds)) + } + start = nil + } + } + } + return result.sorted { $0.startSeconds < $1.startSeconds } + } +} + +// MARK: - Feature Loader + +/// Chunk iterator over a mel feature sequence, mirroring NeMo's `streaming_feat_loader`: +/// fixed core stride, left context of 0 (all preview profiles), right context shrinking at +/// the tail so trailing audio is still emitted. +public struct Nemotron3FeatureLoader { + private let lcMel: Int + private let rcMel: Int + private let coreMel: Int + private let melFeatures: Int + private let capacityMel: Int + + private let featSeq: [Float] + private let featLength: Int + private let featSeqLength: Int + + private var startFeat = 0 + + public init(config: Nemotron3Config, featSeq: [Float], featLength: Int, featSeqLength: Int) { + self.lcMel = config.chunkLeftContext * config.subsamplingFactor + self.rcMel = config.chunkRightContext * config.subsamplingFactor + self.coreMel = config.chunkLen * config.subsamplingFactor + self.melFeatures = config.melFeatures + self.capacityMel = config.chunkMelFrames + self.featSeq = featSeq + self.featLength = featLength + self.featSeqLength = featSeqLength + } + + public mutating func next() -> (features: [Float], length: Int, leftOffset: Int, rightOffset: Int)? { + guard startFeat < featLength else { return nil } + let leftOffset = min(lcMel, startFeat) + let endFeat = min(startFeat + coreMel, featLength) + let rightOffset = min(rcMel, featLength - endFeat) + + let startIdx = (startFeat - leftOffset) * melFeatures + let endIdx = (endFeat + rightOffset) * melFeatures + var features = Array(featSeq[startIdx.. (features: [Float], length: Int, leftOffset: Int, rightOffset: Int)? { + computeMel(final: final) + let sub = config.subsamplingFactor + let lcMel = config.chunkLeftContext * sub + let rcMel = config.chunkRightContext * sub + let coreMel = config.chunkLen * sub + let total = melFramesComputed + let coreStart = nextCoreMel + if final { + guard coreStart < total else { return nil } + } else { + guard coreStart + coreMel + rcMel <= total else { return nil } + } + let leftOffset = min(lcMel, coreStart) + let endFeat = min(coreStart + coreMel, total) + let rightOffset = min(rcMel, total - endFeat) + let frames = endFeat + rightOffset - (coreStart - leftOffset) + + let lo = (coreStart - leftOffset - melCacheStart) * config.melFeatures + var features = Array(melCache[lo..<(lo + frames * config.melFeatures)]) + let capacity = config.chunkMelFrames * config.melFeatures + if features.count < capacity { + features.append(contentsOf: repeatElement(0, count: capacity - features.count)) + } + + nextCoreMel = endFeat + let drop = nextCoreMel - lcMel - melCacheStart + if drop > 0 { + melCache.removeFirst(drop * config.melFeatures) + melCacheStart += drop + } + return (features, frames, leftOffset, rightOffset) + } + + /// Extend the mel cache with every frame whose STFT window is fully available; + /// `final` zero-pads the right edge exactly like center-mode extraction. + private mutating func computeMel(final: Bool) { + let hop = mel.hopLength + let half = mel.nFFT / 2 + let received = samplesReceived + let target: Int + if final { + // Center-padded frame count: 1 + (N + 2*half - win) / hop. + target = received > 0 ? 1 + (received + 2 * half - mel.winLength) / hop : 0 + } else { + target = received >= half ? (received - half) / hop + 1 : 0 + } + let done = melFramesComputed + guard target > done else { return } + + // Frame f is centered on sample f*hop; its window spans ±half around it. Like the + // batch path, pre-emphasize the received samples only, so the zero padding on both + // edges stays zero rather than picking up a -preemph * x[N-1] spike. + let sliceStart = done * hop - half + let sliceEnd = (target - 1) * hop + half + var slice = [Float](repeating: 0, count: sliceEnd - sliceStart) + let copyStart = max(sliceStart, 0) + let copyEnd = min(sliceEnd, received) + if copyEnd > copyStart { + let count = copyEnd - copyStart + let src = copyStart - audioStart + slice.withUnsafeMutableBufferPointer { dst in + audio.withUnsafeBufferPointer { buf in + let out = dst.baseAddress!.advanced(by: copyStart - sliceStart) + let input = buf.baseAddress!.advanced(by: src) + var negPreemph = -preemph + if src > 0 { + // y[n] = x[n] - preemph * x[n-1], seeded from the retained history sample. + vDSP_vsma(input - 1, 1, &negPreemph, input, 1, out, 1, vDSP_Length(count)) + } else { + // Stream start: x[-1] is the zero pad. + out[0] = input[0] + if count > 1 { + vDSP_vsma(input, 1, &negPreemph, input + 1, 1, out + 1, 1, vDSP_Length(count - 1)) + } + } + } + } + } + + let (frames, _, _) = mel.computeFlatTransposed( + audio: slice, paddingMode: .prePadded, expectedFrameCount: target - done) + melCache.append(contentsOf: frames[0..<((target - done) * config.melFeatures)]) + melFramesComputed = target + + // Keep one sample of pre-emphasis history plus the half window behind the next frame. + // In the final flush the last frames' windows extend past the received audio (right + // zero-padding), so the trim point can exceed what is buffered: clamp to `received`. + let keepFrom = min(max(0, target * hop - half - 1), received) + if keepFrom > audioStart { + audio.removeFirst(keepFrom - audioStart) + audioStart = keepFrom + } + } +} diff --git a/Sources/FluidAudio/Diarizer/Nemotron3/Nemotron3Models.swift b/Sources/FluidAudio/Diarizer/Nemotron3/Nemotron3Models.swift new file mode 100644 index 000000000..0669e1d78 --- /dev/null +++ b/Sources/FluidAudio/Diarizer/Nemotron3/Nemotron3Models.swift @@ -0,0 +1,461 @@ +@preconcurrency import CoreML +import Accelerate +import Foundation + +/// Container for the Nemotron 3 Diarization CoreML model and the learned silence embedding. +/// +/// Loaded from a local directory only (NVIDIA evaluation license — the converted models are +/// not distributed). The directory must contain the variant's `.mlmodelc` (or `.mlpackage`) +/// and `learnable_sil_emb.bin` (512 float32, little-endian). +public struct Nemotron3Models { + public let model: MLModel + /// Learned silence embedding, [512]. Used for disabled frames during cache compression. + public let silenceEmbedding: [Float] + /// FeatureStacking projection, [1024, 512] row-major (W^T). Split-graph mode only. + public let preEncodeProjection: [Float]? + public let compilationDuration: TimeInterval + + private let chunkArray: MLMultiArray + private let chunkLengthArray: MLMultiArray + private let spkcacheArray: MLMultiArray + private let spkcacheLengthArray: MLMultiArray + private let fifoArray: MLMultiArray + private let fifoLengthArray: MLMultiArray + private let packedArray: MLMultiArray? + private let attnBiasArray: MLMultiArray? + private let outputMaskArray: MLMultiArray? + private let memoryOptimizer: ANEMemoryOptimizer + + /// Preallocated output backings (fp16, contiguous). Reusing them removes the per-call + /// output IOSurface allocation — the root cause of pool exhaustion on long ANE runs — + /// and shaves per-call marshaling overhead. + private let predsBacking: MLMultiArray + private let hiresBacking: MLMultiArray + private let embsBacking: MLMultiArray? + private let predictionOptions: MLPredictionOptions + + private static let logger = AppLogger(category: "Nemotron3Models") + + public init( + config: Nemotron3Config, + model: MLModel, + silenceEmbedding: [Float], + preEncodeProjection: [Float]? = nil, + compilationDuration: TimeInterval = 0 + ) throws { + self.model = model + self.silenceEmbedding = silenceEmbedding + self.preEncodeProjection = preEncodeProjection + self.compilationDuration = compilationDuration + + self.memoryOptimizer = .init() + self.chunkArray = try memoryOptimizer.createAlignedArray( + shape: [1, NSNumber(value: config.chunkMelFrames), NSNumber(value: config.melFeatures)], + dataType: .float32) + self.spkcacheArray = try memoryOptimizer.createAlignedArray( + shape: [1, NSNumber(value: config.spkcacheLen), NSNumber(value: config.preEncoderDims)], + dataType: .float32) + self.fifoArray = try memoryOptimizer.createAlignedArray( + shape: [1, NSNumber(value: config.fifoLen), NSNumber(value: config.preEncoderDims)], + dataType: .float32) + self.chunkLengthArray = try memoryOptimizer.createAlignedArray(shape: [1], dataType: .int32) + self.spkcacheLengthArray = try memoryOptimizer.createAlignedArray(shape: [1], dataType: .int32) + self.fifoLengthArray = try memoryOptimizer.createAlignedArray(shape: [1], dataType: .int32) + if config.splitGraph { + let t = config.packedFrames + // packed's innermost dim (512) is tile-aligned, so the aligned array is + // contiguous. The mask inputs are NOT: the aligned helper pads innermost + // dims to tile boundaries (output_mask [1,T,1] would get row stride 16, + // scrambling linear writes) — use plain contiguous MLMultiArrays for them. + self.packedArray = try memoryOptimizer.createAlignedArray( + shape: [1, NSNumber(value: t), NSNumber(value: config.preEncoderDims)], + dataType: .float32) + self.attnBiasArray = try MLMultiArray( + shape: [1, 1, 1, NSNumber(value: t)], dataType: .float32) + self.outputMaskArray = try MLMultiArray( + shape: [1, NSNumber(value: t), 1], dataType: .float32) + } else { + self.packedArray = nil + self.attnBiasArray = nil + self.outputMaskArray = nil + } + + // fp16 to match the model's native output precision (CoreML fills them without + // conversion); plain MLMultiArrays are contiguous, which also simplifies readback. + let t = config.packedFrames + let up = config.upsampleFactor + let s = config.numSpeakers + self.predsBacking = try MLMultiArray( + shape: [1, NSNumber(value: t), NSNumber(value: s)], dataType: .float16) + self.hiresBacking = try MLMultiArray( + shape: [1, NSNumber(value: t * up), NSNumber(value: s)], dataType: .float16) + if config.splitGraph { + self.embsBacking = nil + } else { + self.embsBacking = try MLMultiArray( + shape: [1, NSNumber(value: config.chunkEncFrames), NSNumber(value: config.preEncoderDims)], + dataType: .float16) + } + let options = MLPredictionOptions() + var backings: [String: MLMultiArray] = [ + "speaker_preds": predsBacking, + "speaker_preds_10ms": hiresBacking, + ] + if let embsBacking { backings["chunk_pre_encode_embs"] = embsBacking } + options.outputBackings = backings + self.predictionOptions = options + } + + /// Load from a local models directory holding `config.modelFileName` (or its + /// `.mlpackage`) plus `learnable_sil_emb.bin` (and `pre_encode_proj_t.bin` for + /// split-graph presets). + public static func load( + config: Nemotron3Config, + directory: URL, + computeUnits: MLComputeUnits = .all + ) async throws -> Nemotron3Models { + try await load( + config: config, + modelURL: directory.appendingPathComponent(config.modelFileName), + assetsDirectory: directory, + computeUnits: computeUnits) + } + + /// Download the preset's bundle from `FluidInference/nemotron-3-diarization-coreml` + /// (if not cached) and load it. + /// + /// Layout under `cacheDirectory` (default `~/Library/Application Support/FluidAudio/Models`): + /// `nemotron-3-diarization/{monolithic,split}/.mlmodelc` plus the root `.bin` assets. + /// Only the requested preset's bundle is fetched. + public static func loadFromHuggingFace( + config: Nemotron3Config, + cacheDirectory: URL? = nil, + computeUnits: MLComputeUnits = .all, + progressHandler: ProgressHandler? = nil + ) async throws -> Nemotron3Models { + let repo = Repo.nemotron3Diarization + let base = + cacheDirectory + ?? FileManager.default.urls(for: .applicationSupportDirectory, in: .userDomainMask)[0] + .appendingPathComponent("FluidAudio/Models") + let repoDirectory = base.appendingPathComponent(repo.folderName) + let bundlePath = "\(config.hubSubdirectory)/\(config.modelFileName)" + let modelURL = repoDirectory.appendingPathComponent(bundlePath) + let fm = FileManager.default + + // A compiled bundle is complete once its manifest is on disk; partial downloads + // resume file-by-file inside `download(subdirectory:)`. + if !fm.fileExists(atPath: modelURL.appendingPathComponent("coremldata.bin").path) { + logger.info("Downloading \(bundlePath) from \(repo.remotePath)...") + try await ModelHub.download( + repo, subdirectory: bundlePath, to: repoDirectory, progressHandler: progressHandler) + } + + var assets = [ModelNames.Nemotron3.silenceEmbeddingFile] + if config.splitGraph { assets.append(ModelNames.Nemotron3.preEncodeProjectionFile) } + let missing = Set(assets.filter { !fm.fileExists(atPath: repoDirectory.appendingPathComponent($0).path) }) + if !missing.isEmpty { + // Root listing with everything but the wanted files skipped (directories are + // pruned before recursion, so the bundles are never re-listed). + try await ModelHub.download( + repo, subdirectory: "", to: repoDirectory, progressHandler: nil, + shouldSkip: { !missing.contains($0) }) + } + + return try await load( + config: config, modelURL: modelURL, assetsDirectory: repoDirectory, computeUnits: computeUnits) + } + + private static func load( + config: Nemotron3Config, + modelURL requestedModelURL: URL, + assetsDirectory directory: URL, + computeUnits: MLComputeUnits + ) async throws -> Nemotron3Models { + let start = Date() + + var modelURL = requestedModelURL + if !FileManager.default.fileExists(atPath: modelURL.path) { + // Fall back to the uncompiled mlpackage next to the expected mlmodelc. + let packageURL = modelURL.deletingPathExtension().appendingPathExtension("mlpackage") + guard FileManager.default.fileExists(atPath: packageURL.path) else { + throw Nemotron3Error.modelLoadFailed( + "Neither \(config.modelFileName) nor its .mlpackage found at \(modelURL.deletingLastPathComponent().path)" + ) + } + modelURL = try await MLModel.compileModel(at: packageURL) + } + + let mlConfig = MLModelConfiguration() + mlConfig.computeUnits = computeUnits + let model = try MLModel(contentsOf: modelURL, configuration: mlConfig) + + let silURL = directory.appendingPathComponent(ModelNames.Nemotron3.silenceEmbeddingFile) + guard let silData = try? Data(contentsOf: silURL) else { + throw Nemotron3Error.modelLoadFailed("Missing learnable_sil_emb.bin in \(directory.path)") + } + let silCount = silData.count / MemoryLayout.size + guard silCount == config.preEncoderDims else { + throw Nemotron3Error.modelLoadFailed( + "learnable_sil_emb.bin has \(silCount) floats, expected \(config.preEncoderDims)") + } + let silenceEmbedding = silData.withUnsafeBytes { Array($0.bindMemory(to: Float.self)) } + + var projection: [Float]? = nil + if config.splitGraph { + let projURL = directory.appendingPathComponent(ModelNames.Nemotron3.preEncodeProjectionFile) + guard let projData = try? Data(contentsOf: projURL), + projData.count == 1024 * 512 * MemoryLayout.size + else { + throw Nemotron3Error.modelLoadFailed( + "Split-graph mode requires pre_encode_proj_t.bin ([1024,512] fp32) in \(directory.path)") + } + projection = projData.withUnsafeBytes { Array($0.bindMemory(to: Float.self)) } + } + + let duration = Date().timeIntervalSince(start) + logger.info("Loaded Nemotron 3 diarization model in \(String(format: "%.2f", duration))s") + return try Nemotron3Models( + config: config, model: model, silenceEmbedding: silenceEmbedding, + preEncodeProjection: projection, compilationDuration: duration) + } + + // MARK: - Split-graph inference + + /// Run one streaming step through the split graph: host does feature stacking, the + /// 1024->512 projection, state packing, and mask construction; the model is the pure + /// transformer+head. Returns the same `Output` contract (chunk embeddings host-computed). + public func runSplit( + chunk: [Float], + chunkLength: Int, + state: Nemotron3StreamingState, + config: Nemotron3Config + ) throws -> Output { + guard let packedArray, let attnBiasArray, let outputMaskArray, + let projection = preEncodeProjection + else { + throw Nemotron3Error.invalidState("runSplit called on a non-split configuration") + } + let d = config.preEncoderDims + let sub = config.subsamplingFactor + let t = config.packedFrames + + var tStage = Date() + // Feature stacking is a pure reshape of the zero-padded fixed-size mel buffer: + // [mel, 128] row-major == [mel/8, 1024]. Project with one sgemm. + let encCapacity = config.chunkMelFrames / sub + var chunkEmbs = [Float](repeating: 0, count: encCapacity * d) + chunk.withUnsafeBufferPointer { src in + chunkEmbs.withUnsafeMutableBufferPointer { dst in + projection.withUnsafeBufferPointer { proj in + cblas_sgemm( + CblasRowMajor, CblasNoTrans, CblasNoTrans, + Int32(encCapacity), Int32(d), Int32(1024), + 1.0, src.baseAddress, Int32(1024), + proj.baseAddress, Int32(d), + 0.0, dst.baseAddress, Int32(d)) + } + } + } + let encLen = (chunkLength + sub - 1) / sub + + // Pack [spkcache | fifo | chunk] valid frames, zero-pad, build masks. + let packedPtr = packedArray.dataPointer.bindMemory(to: Float.self, capacity: t * d) + var pos = 0 + for (buffer, n) in [ + (state.spkcache, state.spkcacheLength), (state.fifo, state.fifoLength), + (chunkEmbs, encLen), + ] { + buffer.withUnsafeBufferPointer { src in + packedPtr.advanced(by: pos * d).update(from: src.baseAddress!, count: n * d) + } + pos += n + } + if pos < t { + packedPtr.advanced(by: pos * d).update(repeating: 0, count: (t - pos) * d) + } + let biasPtr = attnBiasArray.dataPointer.bindMemory(to: Float.self, capacity: t) + let maskPtr = outputMaskArray.dataPointer.bindMemory(to: Float.self, capacity: t) + biasPtr.update(repeating: 0, count: pos) + biasPtr.advanced(by: pos).update(repeating: -30000.0, count: t - pos) + maskPtr.update(repeating: 1, count: pos) + maskPtr.advanced(by: pos).update(repeating: 0, count: t - pos) + + let inputs = try MLDictionaryFeatureProvider(dictionary: [ + "packed": MLFeatureValue(multiArray: packedArray), + "attn_bias": MLFeatureValue(multiArray: attnBiasArray), + "output_mask": MLFeatureValue(multiArray: outputMaskArray), + ]) + let inputPrepSeconds = Date().timeIntervalSince(tStage) + + tStage = Date() + let output = try model.prediction(from: inputs, options: predictionOptions) + let predictSeconds = Date().timeIntervalSince(tStage) + tStage = Date() + + let predsArray = output.featureValue(for: "speaker_preds")?.multiArrayValue ?? predsBacking + let hiresArray = output.featureValue(for: "speaker_preds_10ms")?.multiArrayValue ?? hiresBacking + return Output( + predictions: Self.floats(from: predsArray), + highResPredictions: Self.floats(from: hiresArray), + chunkEmbeddings: chunkEmbs, + chunkLength: encLen, + inputPrepSeconds: inputPrepSeconds, + predictSeconds: predictSeconds, + readbackSeconds: Date().timeIntervalSince(tStage) + ) + } + + // MARK: - Inference + + public struct Output { + /// 80 ms packed predictions [spkcacheLen + fifoLen + chunkEncFrames, 8] flattened. + public let predictions: [Float] + /// 10 ms packed predictions [(spkcacheLen + fifoLen + chunkEncFrames) * 8, 8] flattened. + public let highResPredictions: [Float] + /// Chunk pre-encode embeddings [chunkEncFrames, 512] flattened. + public let chunkEmbeddings: [Float] + /// Valid encoder frames in `chunkEmbeddings`. + public let chunkLength: Int + /// Per-call wall time split: input tensor copies, CoreML predict, output readback. + public let inputPrepSeconds: Double + public let predictSeconds: Double + public let readbackSeconds: Double + } + + /// Run one streaming step. + /// + /// - Parameters: + /// - chunk: Mel features [chunkMelFrames * 128] flattened (zero-padded to capacity). + /// - chunkLength: Valid mel frames. + /// - state: Current streaming state (read-only here). + public func run( + chunk: [Float], + chunkLength: Int, + state: Nemotron3StreamingState, + config: Nemotron3Config + ) throws -> Output { + var tStage = Date() + memoryOptimizer.optimizedCopy(from: chunk, to: chunkArray, pad: true) + memoryOptimizer.optimizedCopy(from: state.spkcache, to: spkcacheArray, pad: true) + memoryOptimizer.optimizedCopy(from: state.fifo, to: fifoArray, pad: true) + chunkLengthArray[0] = NSNumber(value: Int32(chunkLength)) + spkcacheLengthArray[0] = NSNumber(value: Int32(state.spkcacheLength)) + fifoLengthArray[0] = NSNumber(value: Int32(state.fifoLength)) + + let inputs = try MLDictionaryFeatureProvider(dictionary: [ + "chunk": MLFeatureValue(multiArray: chunkArray), + "chunk_lengths": MLFeatureValue(multiArray: chunkLengthArray), + "spkcache": MLFeatureValue(multiArray: spkcacheArray), + "spkcache_lengths": MLFeatureValue(multiArray: spkcacheLengthArray), + "fifo": MLFeatureValue(multiArray: fifoArray), + "fifo_lengths": MLFeatureValue(multiArray: fifoLengthArray), + ]) + let inputPrepSeconds = Date().timeIntervalSince(tStage) + + tStage = Date() + let output = try model.prediction(from: inputs, options: predictionOptions) + let predictSeconds = Date().timeIntervalSince(tStage) + tStage = Date() + + // Outputs land in the preallocated backings; fall back to the provider's arrays + // if the runtime declined a backing (e.g. shape/dtype mismatch on some OS). + let predsArray = output.featureValue(for: "speaker_preds")?.multiArrayValue ?? predsBacking + let hiresArray = output.featureValue(for: "speaker_preds_10ms")?.multiArrayValue ?? hiresBacking + let embsArray = + output.featureValue(for: "chunk_pre_encode_embs")?.multiArrayValue ?? embsBacking + guard let embsArray else { + throw Nemotron3Error.inferenceFailed("Missing chunk_pre_encode_embs output") + } + let preds = Self.floats(from: predsArray) + let hires = Self.floats(from: hiresArray) + let embs = Self.floats(from: embsArray) + // Advisory only: on GPU-scheduled graphs its fp16 floor_div can be off by one for + // large offline chunks, so derive the valid length host-side instead. + let chunkEncLength = (chunkLength + config.subsamplingFactor - 1) / config.subsamplingFactor + + return Output( + predictions: preds, + highResPredictions: hires, + chunkEmbeddings: embs, + chunkLength: chunkEncLength, + inputPrepSeconds: inputPrepSeconds, + predictSeconds: predictSeconds, + readbackSeconds: Date().timeIntervalSince(tStage) + ) + } + + /// Direct-pointer MLMultiArray -> [Float] copy, honoring strides (reading a strided + /// array through the contiguous fast path scrambles element order — FluidAudio #612). + /// + /// The model's outputs are fp16 with padded rows (e.g. shape [1, T, 8] with row + /// stride 16), so this does one bulk fp16->fp32 conversion over the padded extent + /// followed by a single `vDSP_mmov` 2D compaction, instead of per-element NSNumber + /// reads or `shapedArrayValue` (~1.6 ms/chunk at fast32). + static func floats(from array: MLMultiArray) -> [Float] { + let shape = array.shape.map(\.intValue) + let strides = array.strides.map(\.intValue) + let count = shape.reduce(1, *) + + // Collapse leading singleton dims to a rows x cols view with unit column stride. + // All model outputs are [1, R, C]; also handle fully contiguous arrays as one row. + var rows = 1 + var cols = count + var rowStride = count + if let last = strides.last, last == 1 { + if shape.count >= 2, shape.dropLast(2).allSatisfy({ $0 == 1 }) { + rows = shape[shape.count - 2] + cols = shape[shape.count - 1] + rowStride = strides[strides.count - 2] + } else if strides == (0.. [Float] { + var result = [Float](repeating: 0, count: count) + for i in 0.. Nemotron3ChunkResult { + let d = config.preEncoderDims + let s = config.numSpeakers + let up = config.upsampleFactor + let maxChunk = config.chunkEncFrames - lc - rc + let fifoCap = config.fifoLen + let scCap = config.spkcacheLen + + let scLen = state.spkcacheLength + let fifoLen = state.fifoLength + let chunkLen = min(max(chunkEncLength - lc, 0), maxChunk) + + // Region slices of the packed predictions (valid frames are left-packed). + let spkcachePredsCur = Array(predictions[0..<(scLen * s)]) + let fifoPredsCur = Array(predictions[(scLen * s)..<((scLen + fifoLen) * s)]) + let chunkPredStart = (scLen + fifoLen + lc) * s + let chunkPredsCur = Array(predictions[chunkPredStart..<(chunkPredStart + chunkLen * s)]) + + // High-resolution output for this chunk's core region (NeMo + // `_extract_async_high_resolution_chunk_preds`), taken before the state mutates. + let hiStart = (scLen + fifoLen + lc) * up * s + let hiCount = chunkLen * up * s + let chunkResult = Nemotron3ChunkResult( + probabilities: Array(highResPredictions[hiStart..<(hiStart + hiCount)]), + frameCount: chunkLen * up, + numSpeakers: s + ) + + // FIFO pop lengths (NeMo `_compute_async_fifo_pop_lengths`). + let combined = fifoLen + chunkLen + var pop = 0 + if combined > fifoCap { + pop = min(combined, max(config.spkcacheUpdatePeriod, combined - fifoCap)) + } + if chunkLen == 0 { + pop = fifoLen // finalized stream: flush remaining FIFO into the cache + } + let newFifoLen = combined - pop + + // Logical [FIFO | chunk core] concatenation. + var logicalEmbs = [Float]() + logicalEmbs.reserveCapacity(combined * d) + logicalEmbs.append(contentsOf: state.fifo[0..<(fifoLen * d)]) + logicalEmbs.append(contentsOf: chunkEmbeddings[(lc * d)..<((lc + chunkLen) * d)]) + var logicalPreds = [Float]() + logicalPreds.reserveCapacity(combined * s) + logicalPreds.append(contentsOf: fifoPredsCur) + logicalPreds.append(contentsOf: chunkPredsCur) + + let popEmbs = Array(logicalEmbs[0..<(pop * d)]) + let popPreds = Array(logicalPreds[0..<(pop * s)]) + + // Retained frames become the new FIFO (zero-padded to capacity). + replaceRegion(&state.fifo, with: logicalEmbs[(pop * d)..<(combined * d)], capacity: fifoCap * d) + replaceRegion(&state.fifoPreds, with: logicalPreds[(pop * s)..<(combined * s)], capacity: fifoCap * s) + state.fifoLength = newFifoLen + + // No running silence profile: the checkpoint uses a learned silence embedding. + + // Speaker cache append + compression (NeMo `_update_async_spkcache`). + let updatedLen = scLen + pop + if updatedLen > scCap { + // Candidate preds: fresh predictions for the cache region on the FIRST compression, + // stored (frozen) predictions afterwards. + var candidateEmbs = Array(state.spkcache[0..<(scLen * d)]) + candidateEmbs.append(contentsOf: popEmbs) + var candidatePreds = + state.spkcacheCompressed + ? Array(state.spkcachePreds[0..<(scLen * s)]) + : spkcachePredsCur + candidatePreds.append(contentsOf: popPreds) + + let (newCache, newCachePreds) = compressSpkcache( + embs: candidateEmbs, preds: candidatePreds, frameCount: updatedLen) + replaceRegion(&state.spkcache, with: newCache[...], capacity: scCap * d) + replaceRegion(&state.spkcachePreds, with: newCachePreds[...], capacity: scCap * s) + state.spkcacheLength = scCap + state.spkcacheCompressed = true + } else if pop > 0 { + state.spkcache.replaceSubrange((scLen * d)..<(updatedLen * d), with: popEmbs) + state.spkcachePreds.replaceSubrange((scLen * s)..<(updatedLen * s), with: popPreds) + state.spkcacheLength = updatedLen + } + + return chunkResult + } + + /// Overwrite a fixed-capacity flattened buffer with new content, zero-padding the tail. + private func replaceRegion(_ buffer: inout [Float], with content: ArraySlice, capacity: Int) { + buffer.replaceSubrange(0.. (cache: [Float], cachePreds: [Float]) { + let d = config.preEncoderDims + let s = config.numSpeakers + let scCap = config.spkcacheLen + let silFrames = config.spkcacheSilFramesPerSpk + + let perSpk = scCap / s - silFrames + let strongBoost = Int(Float(perSpk) * config.strongBoostRate) + let weakBoost = Int(Float(perSpk) * config.weakBoostRate) + let minPosScores = Int(Float(perSpk) * config.minPosScoresRate) + + var scores = logPredScores(preds: preds, frameCount: frameCount) + disableLowScores(preds: preds, scores: &scores, frameCount: frameCount, minPosScores: minPosScores) + + // Boost newly added frames (indices beyond the previous cache capacity). + if config.scoresBoostLatest > 0 && frameCount > scCap { + for frame in scCap.. [Float] { + let s = config.numSpeakers + let threshold = config.predScoreThreshold + let count = frameCount * s + var scores = [Float](repeating: 0, count: count) + var logP = [Float](repeating: 0, count: count) + var log1P = [Float](repeating: 0, count: count) + var tmp = [Float](repeating: 0, count: count) + + let p = Array(preds[0..= minPosScores positive-scored frames. + private func disableLowScores( + preds: [Float], scores: inout [Float], frameCount: Int, minPosScores: Int + ) { + let s = config.numSpeakers + var posCounts = [Int](repeating: 0, count: s) + for frame in 0.. 0.5 && scores[i] > 0 { posCounts[spk] += 1 } + } + } + for frame in 0..= minPosScores { + scores[i] = -.infinity + } + } + } + } + + /// NeMo `_boost_topk_scores`: add scaleFactor * log(2) to each speaker's top-k finite scores. + private func boostTopKScores( + scores: inout [Float], frameCount: Int, k: Int, scaleFactor: Float + ) { + let s = config.numSpeakers + guard k > 0, frameCount > 0 else { return } + let delta = scaleFactor * logf(2) + let kEff = min(k, frameCount) + + var topFrames = [Int](repeating: 0, count: kEff) + var topScores = [Float](repeating: -.greatestFiniteMagnitude, count: kEff) + + for spk in 0.. 0 && v > topScores[pos - 1] { + topScores[pos] = topScores[pos - 1] + topFrames[pos] = topFrames[pos - 1] + pos -= 1 + } + topScores[pos] = v + topFrames[pos] = frame + count += 1 + } else { + if v <= topScores[count - 1] { continue } + var pos = count - 1 + while pos > 0 && v > topScores[pos - 1] { + topScores[pos] = topScores[pos - 1] + topFrames[pos] = topFrames[pos - 1] + pos -= 1 + } + topScores[pos] = v + topFrames[pos] = frame + } + } + for i in 0.. (indices: [Int], isDisabled: [Bool]) { + let s = config.numSpeakers + let silFrames = config.spkcacheSilFramesPerSpk + let nFramesNoSil = frameCount - silFrames + let maxIndex = config.maxIndex + let n = frameCount * s + let kEff = min(k, n) + + // Top-k over permuted index space (spk * frameCount + frame), kept DESC by score with + // smaller-index tie-break (matches torch.topk + sort behavior). + var bestIdx = [Int](repeating: 0, count: kEff) + var bestVal = [Float](repeating: -.infinity, count: kEff) + var count = 0 + + for spk in 0.. 0 { + let pv = bestVal[pos - 1] + let pi = bestIdx[pos - 1] + if v > pv || (v == pv && permutedIdx < pi) { + bestVal[pos] = pv + bestIdx[pos] = pi + pos -= 1 + } else { + break + } + } + bestVal[pos] = v + bestIdx[pos] = permutedIdx + count += 1 + } else { + let worstV = bestVal[kEff - 1] + let worstI = bestIdx[kEff - 1] + if v < worstV || (v == worstV && permutedIdx >= worstI) { continue } + var pos = kEff - 1 + while pos > 0 { + let pv = bestVal[pos - 1] + let pi = bestIdx[pos - 1] + if v > pv || (v == pv && permutedIdx < pi) { + bestVal[pos] = pv + bestIdx[pos] = pi + pos -= 1 + } else { + break + } + } + bestVal[pos] = v + bestIdx[pos] = permutedIdx + } + } + } + + var topK = [Int](repeating: maxIndex, count: k) + for i in 0..= nFramesNoSil { + isDisabled[i] = true + topK[i] = 0 + } + } + return (topK, isDisabled) + } +} diff --git a/Sources/FluidAudio/Diarizer/Nemotron3/Nemotron3Types.swift b/Sources/FluidAudio/Diarizer/Nemotron3/Nemotron3Types.swift new file mode 100644 index 000000000..f057bcaa4 --- /dev/null +++ b/Sources/FluidAudio/Diarizer/Nemotron3/Nemotron3Types.swift @@ -0,0 +1,278 @@ +import Foundation + +// MARK: - Configuration + +/// Configuration for Nemotron 3 Diarization streaming inference (8-speaker streaming Sortformer). +/// +/// Mirrors NeMo `SortformerModules` parameters for `nvidia/Nemotron-3-Diarization-preview`. +/// Latency = (chunkLen + chunkRightContext) * 80 ms. +/// +/// - Important: The preview checkpoint is under an NVIDIA evaluation license. Converted CoreML +/// models are loaded from a local directory only — there is no HuggingFace download path. +public struct Nemotron3Config: Sendable { + + // MARK: Architecture (fixed by the checkpoint) + + public let numSpeakers: Int = 8 + public let preEncoderDims: Int = 512 + public let subsamplingFactor: Int = 8 + public let melFeatures: Int = 128 + public let sampleRate: Int = 16000 + + /// High-resolution output upsample factor (80 ms encoder frame -> 10 ms output frames). + public let upsampleFactor: Int = 8 + + // MARK: Streaming parameters (must match the converted model's fixed shapes) + + public var chunkLen: Int + public var chunkLeftContext: Int + public var chunkRightContext: Int + public var fifoLen: Int + public var spkcacheLen: Int + public var spkcacheUpdatePeriod: Int + + // MARK: Compression constants (NeMo model_config.yaml) + + public var silenceThreshold: Float = 0.2 + public var predScoreThreshold: Float = 0.25 + public var scoresBoostLatest: Float = 0.05 + public var strongBoostRate: Float = 0.75 + public var weakBoostRate: Float = 1.5 + public var minPosScoresRate: Float = 0.5 + public var spkcacheSilFramesPerSpk: Int = 1 + public let maxIndex: Int = 99999 + + public var debugMode: Bool = false + + /// Model file name inside the models directory, e.g. `Nemotron3Diarizer_low.mlmodelc`. + public var modelFileName: String + + /// Split-graph mode: the model contains only the pure-fp transformer+head + /// (inputs `packed`/`attn_bias`/`output_mask`); feature stacking, the 1024->512 + /// projection, state packing, and mask construction run host-side. Requires + /// `pre_encode_proj_t.bin` next to the model. Runs 100% ANE-resident and is not + /// subject to the monolithic graph's chunk-length ANECCompile cliff. + public var splitGraph: Bool = false + + // MARK: Derived + + /// Mel frames the CoreML `chunk` input expects: (lc + chunk + rc) * 8. + public var chunkMelFrames: Int { + (chunkLeftContext + chunkLen + chunkRightContext) * subsamplingFactor + } + + /// Encoder frames of the chunk region (physical capacity incl. contexts). + public var chunkEncFrames: Int { + chunkLeftContext + chunkLen + chunkRightContext + } + + /// Packed sequence length of the model output: spkcache + fifo + chunk regions. + public var packedFrames: Int { + spkcacheLen + fifoLen + chunkEncFrames + } + + /// Output frame duration for high-resolution predictions (10 ms). + public var outputFrameSeconds: Float { 0.01 } + + /// Seconds of audio consumed per streaming step (the chunk core). + public var chunkSeconds: Double { + Double(chunkLen * subsamplingFactor) * 0.01 + } + + /// Input-buffer latency in seconds: core + right context. + public var latencySeconds: Double { + Double((chunkLen + chunkRightContext) * subsamplingFactor) * 0.01 + } + + /// Subdirectory of `FluidInference/nemotron-3-diarization-coreml` holding this preset's bundle. + public var hubSubdirectory: String { splitGraph ? "split" : "monolithic" } + + // MARK: Presets (model card recommended profiles) + + /// 30.4 s input-buffer latency, offline-style quality; highest-throughput batch profile. + public static let offline = Nemotron3Config( + chunkLen: 340, chunkRightContext: 40, fifoLen: 40, spkcacheUpdatePeriod: 300, + modelFileName: "Nemotron3Diarizer_offline.mlmodelc") + + /// 1.04 s latency streaming. + public static let low = Nemotron3Config( + chunkLen: 9, chunkRightContext: 4, fifoLen: 264, spkcacheUpdatePeriod: 222, + modelFileName: "Nemotron3Diarizer_low.mlmodelc") + + /// 0.64 s latency streaming. + public static let veryLow = Nemotron3Config( + chunkLen: 6, chunkRightContext: 2, fifoLen: 264, spkcacheUpdatePeriod: 222, + modelFileName: "Nemotron3Diarizer_verylow.mlmodelc") + + /// 0.32 s latency streaming. + public static let ultraLow = Nemotron3Config( + chunkLen: 3, chunkRightContext: 1, fifoLen: 264, spkcacheUpdatePeriod: 222, + modelFileName: "Nemotron3Diarizer_ultra.mlmodelc") + + /// 1.04 s latency with a 40-frame FIFO: packed sequence 317 vs 541 frames — + /// substantially faster per call on ANE than `low` at a small quality cost. + /// Not a model-card profile. + public static let fast = Nemotron3Config( + chunkLen: 9, chunkRightContext: 4, fifoLen: 40, spkcacheUpdatePeriod: 40, + modelFileName: "Nemotron3Diarizer_fast.mlmodelc") + + /// 4.16 s latency, 48-frame chunk: amortizes the static spkcache+FIFO cost per call for + /// high-throughput batch/near-live use. + public static let efficient = Nemotron3Config( + chunkLen: 48, chunkRightContext: 4, fifoLen: 264, spkcacheUpdatePeriod: 222, + modelFileName: "Nemotron3Diarizer_efficient.mlmodelc") + + /// 2.24 s latency, 1.92 s audio per call at `fast`-class per-call cost — + /// bigger chunks recover the small-FIFO quality penalty. + public static let fast24 = Nemotron3Config( + chunkLen: 24, chunkRightContext: 4, fifoLen: 40, spkcacheUpdatePeriod: 40, + modelFileName: "Nemotron3Diarizer_fast24.mlmodelc") + + /// 2.88 s latency, 2.56 s audio per call — matches the card-standard `low` + /// profile's quality at a fraction of its per-call ANE cost. Recommended default + /// when latency up to ~3 s is acceptable. + public static let fast32 = Nemotron3Config( + chunkLen: 32, chunkRightContext: 4, fifoLen: 40, spkcacheUpdatePeriod: 40, + modelFileName: "Nemotron3Diarizer_fast32.mlmodelc") + + /// 10.56 s latency, 10.24 s audio per call; largest monolithic chunk that still + /// compiles for ANE (192 fails ANECCompile). Best quality of the streaming preset + /// lineup. High-throughput near-live tier; for pure GPU batch prefer `.offline`. + public static let fast128 = Nemotron3Config( + chunkLen: 128, chunkRightContext: 4, fifoLen: 40, spkcacheUpdatePeriod: 40, + modelFileName: "Nemotron3Diarizer_fast128.mlmodelc") + + public init( + chunkLen: Int, + chunkLeftContext: Int = 0, + chunkRightContext: Int, + fifoLen: Int, + spkcacheLen: Int = 264, + spkcacheUpdatePeriod: Int, + modelFileName: String, + splitGraph: Bool = false + ) { + self.chunkLen = chunkLen + self.chunkLeftContext = chunkLeftContext + self.chunkRightContext = chunkRightContext + self.fifoLen = fifoLen + self.spkcacheLen = spkcacheLen + self.spkcacheUpdatePeriod = spkcacheUpdatePeriod + self.modelFileName = modelFileName + self.splitGraph = splitGraph + } + + public static func preset(named name: String) -> Nemotron3Config? { + // "-int8" selects the int8-quantized model file with identical parameters. + if name.hasSuffix("-int8"), var base = preset(named: String(name.dropLast(5))) { + base.modelFileName = base.modelFileName.replacingOccurrences( + of: ".mlmodelc", with: "_int8.mlmodelc") + return base + } + // "-split" selects a split-graph model (see `splitGraph`); the underlying + // model files use the sweep naming (s32 = fast32's shape). + switch name { + case "offline": return .offline + case "low": return .low + case "verylow": return .veryLow + case "ultra": return .ultraLow + case "fast": return .fast + case "fast24": return .fast24 + case "fast32": return .fast32 + case "fast128": return .fast128 + case "efficient": return .efficient + case "fast32-split": + return Nemotron3Config( + chunkLen: 32, chunkRightContext: 4, fifoLen: 40, spkcacheUpdatePeriod: 40, + modelFileName: "Nemotron3Diarizer_s32_split.mlmodelc", splitGraph: true) + case "fast32-split-w8a8": + return Nemotron3Config( + chunkLen: 32, chunkRightContext: 4, fifoLen: 40, spkcacheUpdatePeriod: 40, + modelFileName: "Nemotron3Diarizer_s32_split_w8a8.mlmodelc", splitGraph: true) + case "c128-split": + return Nemotron3Config( + chunkLen: 128, chunkRightContext: 4, fifoLen: 40, spkcacheUpdatePeriod: 40, + modelFileName: "Nemotron3Diarizer_c128_split.mlmodelc", splitGraph: true) + case "c128-split-w8a8": + return Nemotron3Config( + chunkLen: 128, chunkRightContext: 4, fifoLen: 40, spkcacheUpdatePeriod: 40, + modelFileName: "Nemotron3Diarizer_c128_split_w8a8.mlmodelc", splitGraph: true) + case "c192-split": + return Nemotron3Config( + chunkLen: 192, chunkRightContext: 4, fifoLen: 40, spkcacheUpdatePeriod: 40, + modelFileName: "Nemotron3Diarizer_c192_split.mlmodelc", splitGraph: true) + case "c256-split": + return Nemotron3Config( + chunkLen: 256, chunkRightContext: 4, fifoLen: 40, spkcacheUpdatePeriod: 40, + modelFileName: "Nemotron3Diarizer_c256_split.mlmodelc", splitGraph: true) + default: return nil + } + } +} + +// MARK: - Streaming State + +/// Fixed-capacity streaming state, mirroring NeMo's async `StreamingSortformerState` at batch 1. +/// +/// `spkcache`/`fifo` are always full physical capacity (zero-padded past the valid length), +/// matching the CoreML model's fixed input shapes. +public struct Nemotron3StreamingState: Sendable { + /// [spkcacheLen, 512] flattened, valid frames left-packed. + public var spkcache: [Float] + public var spkcacheLength: Int + /// [spkcacheLen, 8] flattened. Meaningful only from the first compression onward. + public var spkcachePreds: [Float] + public var spkcacheCompressed: Bool + + /// [fifoLen, 512] flattened, valid frames left-packed. + public var fifo: [Float] + public var fifoLength: Int + /// [fifoLen, 8] flattened. + public var fifoPreds: [Float] + + public init(config: Nemotron3Config) { + let d = config.preEncoderDims + let s = config.numSpeakers + self.spkcache = [Float](repeating: 0, count: config.spkcacheLen * d) + self.spkcachePreds = [Float](repeating: 0, count: config.spkcacheLen * s) + self.spkcacheLength = 0 + self.spkcacheCompressed = false + self.fifo = [Float](repeating: 0, count: config.fifoLen * d) + self.fifoPreds = [Float](repeating: 0, count: config.fifoLen * s) + self.fifoLength = 0 + } +} + +// MARK: - Results + +/// Per-chunk streaming result at 10 ms resolution. +public struct Nemotron3ChunkResult: Sendable { + /// Speaker activity probabilities for this chunk's core frames, [frames * 8] flattened, + /// 10 ms per frame. + public let probabilities: [Float] + public let frameCount: Int + public let numSpeakers: Int +} + +/// A contiguous speech segment attributed to one speaker slot (arrival-ordered). +public struct Nemotron3Segment: Sendable { + public let speakerIndex: Int + public let startSeconds: Float + public let endSeconds: Float +} + +// MARK: - Errors + +public enum Nemotron3Error: Error, LocalizedError { + case modelLoadFailed(String) + case inferenceFailed(String) + case invalidState(String) + + public var errorDescription: String? { + switch self { + case .modelLoadFailed(let m): return "Failed to load Nemotron 3 diarization model: \(m)" + case .inferenceFailed(let m): return "Nemotron 3 diarization inference failed: \(m)" + case .invalidState(let m): return "Invalid Nemotron 3 diarization state: \(m)" + } + } +} diff --git a/Sources/FluidAudio/ModelNames.swift b/Sources/FluidAudio/ModelNames.swift index e2085d59f..d8138bb1a 100644 --- a/Sources/FluidAudio/ModelNames.swift +++ b/Sources/FluidAudio/ModelNames.swift @@ -53,6 +53,7 @@ public enum Repo: String, CaseIterable, Sendable { case kokoroAneZh = "FluidInference/kokoro-82m-coreml/ANE-zh" case kokoroAneJa = "FluidInference/kokoro-82m-coreml/ANE-ja" case sortformer = "FluidInference/diar-streaming-sortformer-coreml" + case nemotron3Diarization = "FluidInference/nemotron-3-diarization-coreml" case lseendAmi = "FluidInference/ls-eend-coreml/optimized/ami" case lseendCallHome = "FluidInference/ls-eend-coreml/optimized/ch" case lseendDihard2 = "FluidInference/ls-eend-coreml/optimized/dih2" @@ -174,6 +175,8 @@ public enum Repo: String, CaseIterable, Sendable { return "kokoro-82m-coreml/ANE-ja" case .sortformer: return "diar-streaming-sortformer-coreml" + case .nemotron3Diarization: + return "nemotron-3-diarization-coreml" case .lseendAmi: return "ls-eend-coreml/optimized/ami" case .lseendCallHome: @@ -815,6 +818,14 @@ public enum ModelNames { } /// Sortformer streaming diarization model names + public enum Nemotron3 { + /// Root-level assets every preset needs (split-graph presets also need + /// `pre_encode_proj_t.bin`). + public static let silenceEmbeddingFile = "learnable_sil_emb.bin" + public static let preEncodeProjectionFile = "pre_encode_proj_t.bin" + public static let requiredAssets: Set = [silenceEmbeddingFile, preEncodeProjectionFile] + } + public enum Sortformer { /// Selects which weight-precision build of the model set to download. /// @@ -1712,6 +1723,14 @@ public enum ModelNames { return [variant] } return ModelNames.Sortformer.requiredModels + case .nemotron3Diarization: + // Downloads are driven by `Nemotron3Models.loadFromHuggingFace` via + // `download(subdirectory:)` (one preset bundle + the root .bin assets); + // provided for exhaustiveness. + if let variant = variant { + return [variant] + } + return ModelNames.Nemotron3.requiredAssets case .lseendAmi, .lseendCallHome, .lseendDihard2, .lseendDihard3: if let variant = variant { return [variant + ".mlmodelc"] diff --git a/Sources/FluidAudio/Shared/AudioMelSpectrogram.swift b/Sources/FluidAudio/Shared/AudioMelSpectrogram.swift index dd629bc40..1b6ce956e 100644 --- a/Sources/FluidAudio/Shared/AudioMelSpectrogram.swift +++ b/Sources/FluidAudio/Shared/AudioMelSpectrogram.swift @@ -29,8 +29,8 @@ public final class AudioMelSpectrogram { // Config private let sampleRate: Int public let nFFT: Int - private let hopLength: Int // window_stride * sample_rate - private let winLength: Int // window_size * sample_rate + public let hopLength: Int // window_stride * sample_rate + public let winLength: Int // window_size * sample_rate private let fMin: Float = 0.0 private let fMax: Float // sample_rate / 2 internal let preemph: Float // NeMo preemphasis coefficient @@ -56,13 +56,16 @@ public final class AudioMelSpectrogram { private var imagSq: [Float] private var frame: [Float] + /// NeMo's default pre-emphasis coefficient. + public static let defaultPreemph: Float = 0.97 + public init( sampleRate: Int = 16000, nMels: Int = 128, nFFT: Int = 512, hopLength: Int = 160, winLength: Int = 400, - preemph: Float = 0.97, + preemph: Float = AudioMelSpectrogram.defaultPreemph, padTo: Int = 0, logFloor: Float = powf(2, -24), logFloorMode: LogFloorMode = .additive, diff --git a/Sources/FluidAudioCLI/Commands/DiarizationBenchmarkUtils.swift b/Sources/FluidAudioCLI/Commands/DiarizationBenchmarkUtils.swift index 7bab553b5..2340c5106 100644 --- a/Sources/FluidAudioCLI/Commands/DiarizationBenchmarkUtils.swift +++ b/Sources/FluidAudioCLI/Commands/DiarizationBenchmarkUtils.swift @@ -12,8 +12,20 @@ enum DiarizationBenchmarkUtils { } /// Dataset corpora supported by diarization benchmarks. + /// + /// Card-protocol conditions (NVIDIA Nemotron 3 model card): + /// `ami` = AMI test MHM (Mix-Headset), `amiSdm` = AMI test SDM (Array1-01), + /// `alimeetingFar` = channel 0 of the far-field array, `alimeetingNear` = + /// equal-weight mix of per-speaker headset channels (both pre-materialized by + /// `Scripts/materialize_alimeeting_card_audio.py`), `notsofarMhm`/`notsofarSc` = + /// NOTSOFAR1 eval mixes (see `Scripts/materialize_notsofar_card_audio.py`). enum Dataset: String { case ami = "ami" + case amiSdm = "ami-sdm" + case alimeetingFar = "alimeeting-far" + case alimeetingNear = "alimeeting-near" + case notsofarMhm = "notsofar-mhm" + case notsofarSc = "notsofar-sc" case voxconverse = "voxconverse" case callhome = "callhome" } @@ -36,12 +48,12 @@ enum DiarizationBenchmarkUtils { // MARK: - File Paths - static func getAMIFiles(split: AMISplit = .test, maxFiles: Int?) -> [String] { + static func getAMIFiles(split: AMISplit = .test, dataset: Dataset = .ami, maxFiles: Int?) -> [String] { let allMeetings = getAMIMeetings(split: split) var availableMeetings: [String] = [] for meeting in DatasetDownloader.officialAMITestSet { - let path = getAudioPath(for: meeting, dataset: .ami) + let path = getAudioPath(for: meeting, dataset: dataset) if FileManager.default.fileExists(atPath: path) { availableMeetings.append(meeting) } @@ -53,6 +65,36 @@ enum DiarizationBenchmarkUtils { return availableMeetings } + /// Enumerates meetings for datasets whose audio lives in a flat directory of + /// `.wav` files with a matching reference RTTM per meeting. + static func getDirectoryFiles(dataset: Dataset, maxFiles: Int?) -> [String] { + let sampleAudioPath = getAudioPath(for: "PROBE", dataset: dataset) + let audioDir = URL(fileURLWithPath: sampleAudioPath).deletingLastPathComponent() + + guard + let files = try? FileManager.default.contentsOfDirectory( + at: audioDir, includingPropertiesForKeys: nil) + else { + return [] + } + + var availableMeetings: [String] = [] + for file in files where file.pathExtension == "wav" { + let name = file.deletingPathExtension().lastPathComponent + if let rttmURL = getRTTMURL(for: name, dataset: dataset), + FileManager.default.fileExists(atPath: rttmURL.path) + { + availableMeetings.append(name) + } + } + + availableMeetings.sort() + if let max = maxFiles { + return Array(availableMeetings.prefix(max)) + } + return availableMeetings + } + static func getAMIMeetings(split: AMISplit) -> [String] { switch split { case .train: @@ -102,6 +144,26 @@ enum DiarizationBenchmarkUtils { return homeDir.appendingPathComponent( "FluidAudioDatasets/ami_official/sdm/\(meeting).Mix-Headset.wav" ).path + case .amiSdm: + return homeDir.appendingPathComponent( + "FluidAudioDatasets/ami_official/sdm_true/\(meeting).Array1-01.wav" + ).path + case .alimeetingFar: + return homeDir.appendingPathComponent( + "FluidAudioDatasets/alimeeting/card/far_ch0/\(meeting).wav" + ).path + case .alimeetingNear: + return homeDir.appendingPathComponent( + "FluidAudioDatasets/alimeeting/card/near_mix/\(meeting).wav" + ).path + case .notsofarMhm: + return homeDir.appendingPathComponent( + "FluidAudioDatasets/notsofar/card/eval_mhm/\(meeting).wav" + ).path + case .notsofarSc: + return homeDir.appendingPathComponent( + "FluidAudioDatasets/notsofar/card/eval_sc/\(meeting).wav" + ).path case .voxconverse: return homeDir.appendingPathComponent( "FluidAudioDatasets/voxconverse/voxconverse_test_wav/\(meeting).wav" @@ -114,16 +176,27 @@ enum DiarizationBenchmarkUtils { } static func getRTTMURL(for meeting: String, dataset: Dataset) -> URL? { + let homeDir = FileManager.default.homeDirectoryForCurrentUser switch dataset { - case .ami: + case .ami, .amiSdm: + // MHM and SDM share the same forced-alignment references; only the + // audio condition differs. return getAMIRTTMURL(for: meeting) + case .alimeetingFar, .alimeetingNear: + // nttcslab-sp/diar-forced-alignment publishes one reference set per + // meeting (Test_Ali_far); Near/Far are the same meetings. + return homeDir.appendingPathComponent( + "FluidAudioDatasets/diar-forced-alignment/AliMeeting/Test_Ali_far/\(meeting).rttm" + ) + case .notsofarMhm, .notsofarSc: + return homeDir.appendingPathComponent( + "FluidAudioDatasets/notsofar/card/rttm/\(meeting).rttm" + ) case .voxconverse: - let homeDir = FileManager.default.homeDirectoryForCurrentUser return homeDir.appendingPathComponent( "FluidAudioDatasets/voxconverse/rttm_repo/test/\(meeting).rttm" ) case .callhome: - let homeDir = FileManager.default.homeDirectoryForCurrentUser return homeDir.appendingPathComponent( "FluidAudioDatasets/callhome_eng/rttm/\(meeting).rttm" ) @@ -227,8 +300,10 @@ enum DiarizationBenchmarkUtils { /// Returns files for the given dataset, filtering by availability. static func getFiles(for dataset: Dataset, maxFiles: Int?) -> [String] { switch dataset { - case .ami: - return getAMIFiles(maxFiles: maxFiles) + case .ami, .amiSdm: + return getAMIFiles(dataset: dataset, maxFiles: maxFiles) + case .alimeetingFar, .alimeetingNear, .notsofarMhm, .notsofarSc: + return getDirectoryFiles(dataset: dataset, maxFiles: maxFiles) case .voxconverse: return getVoxConverseFiles(maxFiles: maxFiles) case .callhome: @@ -238,6 +313,18 @@ enum DiarizationBenchmarkUtils { // MARK: - Summary & Output + /// Speaker-counting metrics as defined on the NVIDIA Nemotron 3 model card: + /// SCA = percentage of files where predicted speaker count equals ground truth; + /// MAE = mean of `|predicted - ground truth|` over files. + static func speakerCountMetrics(results: [BenchmarkResult]) -> (sca: Double, mae: Double) { + guard !results.isEmpty else { return (0, 0) } + let exact = results.filter { $0.detectedSpeakers == $0.groundTruthSpeakers }.count + let absErrors = results.map { abs($0.detectedSpeakers - $0.groundTruthSpeakers) } + let sca = Double(exact) / Double(results.count) * 100 + let mae = Double(absErrors.reduce(0, +)) / Double(results.count) + return (sca, mae) + } + /// Prints a formatted benchmark summary table. /// /// - Parameters: diff --git a/Sources/FluidAudioCLI/Commands/Nemotron3DiarizeCommand.swift b/Sources/FluidAudioCLI/Commands/Nemotron3DiarizeCommand.swift new file mode 100644 index 000000000..2109e03ba --- /dev/null +++ b/Sources/FluidAudioCLI/Commands/Nemotron3DiarizeCommand.swift @@ -0,0 +1,705 @@ +#if os(macOS) +import CoreML +import FluidAudio +import Foundation + +/// CLI for Nemotron 3 Diarization. +enum Nemotron3DiarizeCommand { + private static let logger = AppLogger(category: "Nemotron3CLI") + + static func printUsage() { + let usage = """ + Nemotron 3 Diarization (8-speaker streaming) + + Usage: + fluidaudiocli nemotron3-diarize [options] + fluidaudiocli nemotron3-benchmark [options] + + Shared options: + --models Local directory containing Nemotron3Diarizer_.mlmodelc + and learnable_sil_emb.bin. Omit to download the preset from + FluidInference/nemotron-3-diarization-coreml. + --variant offline | low | fast | fast32 | fast128 | fast32-split-w8a8 | + c128-split-w8a8 (default: low; verylow/ultra need --models) + --threshold Speaker activity threshold (default: 0.5) + + nemotron3-diarize options: + --streaming Feed audio in 100 ms pieces through the live path + (appendAudio / processBufferedAudio / finishStream) + --dump-preds Write raw frame probabilities (float32 LE, [T, 8]) for parity checks + --output Write RTTM hypothesis + + nemotron3-benchmark options: + --dataset ami | ami-sdm | alimeeting-far | alimeeting-near | + notsofar-mhm | notsofar-sc | voxconverse | callhome + (default: ami — the AMI test MHM/Mix-Headset condition) + --single-file Process one meeting (e.g. ES2004a) + --max-files Limit number of files + --collar DER collar (default: 0) + --output Output JSON results + """ + fputs(usage, stderr) + fflush(stderr) + } + + struct CustomShape { + var chunkLen: Int? + var rightContext: Int? + var fifoLen: Int? + var spkcacheLen: Int? + var updatePeriod: Int? + } + + private static func loadDiarizer( + modelsDir: String?, variantName: String, custom: CustomShape = CustomShape(), + computeUnits: MLComputeUnits = .all + ) async throws -> (Nemotron3Diarizer, TimeInterval) { + var config: Nemotron3Config + if let preset = Nemotron3Config.preset(named: variantName) { + config = preset + } else { + // Sweep variant: derive shape from flags, model file from the variant name. + config = Nemotron3Config( + chunkLen: custom.chunkLen ?? 9, + chunkRightContext: custom.rightContext ?? 4, + fifoLen: custom.fifoLen ?? 40, + spkcacheLen: custom.spkcacheLen ?? 264, + spkcacheUpdatePeriod: custom.updatePeriod ?? 40, + modelFileName: "Nemotron3Diarizer_\(variantName).mlmodelc") + } + let start = Date() + let models: Nemotron3Models + if let modelsDir { + models = try await Nemotron3Models.load( + config: config, + directory: URL(fileURLWithPath: modelsDir), + computeUnits: computeUnits + ) + } else { + models = try await Nemotron3Models.loadFromHuggingFace(config: config, computeUnits: computeUnits) + } + return (Nemotron3Diarizer(config: config, models: models), Date().timeIntervalSince(start)) + } + + /// Drive the live path with fixed-size audio pieces; returns the same flat + /// `[frames * 8]` probabilities as `processComplete`. + static func streamAudio( + _ audio: [Float], through diarizer: Nemotron3Diarizer, pieceSamples: Int = 1600 + ) throws -> (probabilities: [Float], frameCount: Int) { + diarizer.reset() + var probs: [Float] = [] + var frames = 0 + var offset = 0 + while offset < audio.count { + let end = min(offset + pieceSamples, audio.count) + diarizer.appendAudio(Array(audio[offset.. MLComputeUnits { + switch s { + case "ane": return .cpuAndNeuralEngine + case "gpu": return .cpuAndGPU + case "cpu": return .cpuOnly + default: return .all + } + } + + /// Silero VAD -> per-10ms-frame speech mask, with speech regions padded by + /// `padSeconds` on both sides to protect onsets/offsets at chunk granularity. + /// Pass a shared `VadManager` when calling repeatedly — a fresh instance per file + /// leaks IOSurfaces across a long benchmark run and eventually fails allocation. + static func speechMask( + audio: [Float], threshold: Float, padSeconds: Double = 1.0, + vad existingVad: VadManager? = nil + ) async throws -> [Bool] { + let vad: VadManager + if let existingVad { + vad = existingVad + } else { + vad = try await VadManager(config: VadConfig(defaultThreshold: threshold)) + } + let framesPerVadChunk = VadManager.chunkSize / 160 // 4096 samples -> 25.6 x 10ms frames + let frameCount = (audio.count + 159) / 160 + var mask = [Bool](repeating: false, count: frameCount) + // Process in bounded segments: one monolithic process() over a long meeting churns + // thousands of MLMultiArrays without an autorelease drain and exhausts IOSurfaces + // (same failure class as issue #752). Silero state resets per call anyway; the 1 s + // padding below absorbs boundary effects. + let segmentSamples = 300 * 16000 + var segmentStart = 0 + while segmentStart < audio.count { + let segmentEnd = min(segmentStart + segmentSamples, audio.count) + let results = try await vad.process(Array(audio[segmentStart..= threshold { + let start = frameOffset + i * framesPerVadChunk + let end = min(frameOffset + (i + 1) * framesPerVadChunk + 1, frameCount) + if start < frameCount { + for f in start.. 0 ? s / p.totalSeconds * 100 : 0 + let padded = name.padding(toLength: 14, withPad: " ", startingAt: 0) + print( + padded + + String(format: "%8.3fs %5.1f%% %8.3fms/chunk", s, pct, s / n * 1000)) + } + print("Pipeline profile (\(p.chunkCount) chunks):") + line("mel", p.melSeconds) + line("chunk-slice", p.chunkSliceSeconds) + line("input-prep", p.inputPrepSeconds) + line("predict", p.predictSeconds) + line("readback", p.readbackSeconds) + line("state-update", p.stateUpdateSeconds) + line("output-append", p.outputAppendSeconds) + line("total", p.totalSeconds) + } + + let segments = Nemotron3Diarizer.segments( + probabilities: probs, frameCount: frames, threshold: threshold) + let speakers = Set(segments.map(\.speakerIndex)) + print("Detected \(speakers.count) speakers, \(segments.count) segments") + for seg in segments.prefix(20) { + print( + " spk\(seg.speakerIndex): \(String(format: "%7.2f", seg.startSeconds))s - " + + "\(String(format: "%7.2f", seg.endSeconds))s") + } + if segments.count > 20 { print(" ... (\(segments.count - 20) more)") } + + if let dumpPredsPath { + var data = Data(capacity: probs.count * 4) + probs.withUnsafeBytes { data.append(contentsOf: $0) } + try data.write(to: URL(fileURLWithPath: dumpPredsPath)) + print("Dumped \(frames)x8 frame probabilities to \(dumpPredsPath)") + } + + if let outputPath { + let fileId = URL(fileURLWithPath: audioPath).deletingPathExtension().lastPathComponent + var rttm = "" + for seg in segments { + let dur = seg.endSeconds - seg.startSeconds + rttm += + "SPEAKER \(fileId) 1 \(String(format: "%.3f", seg.startSeconds)) " + + "\(String(format: "%.3f", dur)) speaker_\(seg.speakerIndex) \n" + } + try rttm.write(toFile: outputPath, atomically: true, encoding: .utf8) + print("Wrote RTTM to \(outputPath)") + } + } catch { + print("Error: \(error)") + exit(1) + } + } + + // MARK: - Batch mode (concurrent GPU streams) + + /// Process many files with N concurrent workers, each owning its own model instance. + /// Multi-stream measurement: the M5 Pro GPU takes exactly one extra concurrent stream + /// (+43% aggregate) before saturating, so the default is 2 workers on the GPU route. + static func runBatch(arguments: [String]) async { + var modelsDir: String? + var variantName = "fast32" + var workers = 2 + var computeUnits: MLComputeUnits = .cpuAndGPU + var files: [String] = [] + var dataset: DiarizationBenchmarkUtils.Dataset = .ami + var threshold: Float = 0.5 + var collar: Double = 0 + var maxFiles: Int? + + var i = 0 + while i < arguments.count { + switch arguments[i] { + case "--models": + i += 1 + modelsDir = arguments[safe: i] + case "--variant": + i += 1 + variantName = arguments[safe: i] ?? "fast32" + case "--workers": + i += 1 + workers = arguments[safe: i].flatMap(Int.init) ?? 2 + case "--compute-units": + i += 1 + computeUnits = parseComputeUnits(arguments[safe: i]) + case "--dataset": + i += 1 + dataset = DiarizationBenchmarkUtils.Dataset(rawValue: arguments[safe: i] ?? "ami") ?? .ami + case "--files": + i += 1 + files = arguments[safe: i]?.split(separator: ",").map(String.init) ?? [] + case "--max-files": + i += 1 + maxFiles = arguments[safe: i].flatMap(Int.init) + case "--threshold": + i += 1 + threshold = arguments[safe: i].flatMap(Float.init) ?? 0.5 + case "--collar": + i += 1 + collar = arguments[safe: i].flatMap(Double.init) ?? 0 + case "--help", "-h": + printUsage() + return + default: + break + } + i += 1 + } + + if files.isEmpty { + files = DiarizationBenchmarkUtils.getFiles(for: dataset, maxFiles: maxFiles) + } + guard !files.isEmpty else { + print("No files to process.") + exit(1) + } + + print("Batch: \(files.count) files, \(workers) workers, variant \(variantName)") + let wallStart = Date() + + // Stride-assign files to workers; each worker loads its own model instance so + // CoreML queues the streams independently. + let assignments = (0.. DiarizationBenchmarkUtils.BenchmarkResult? { + let audioPath = DiarizationBenchmarkUtils.getAudioPath(for: meeting, dataset: dataset) + guard FileManager.default.fileExists(atPath: audioPath) else { + print(" Audio not found: \(audioPath)") + return nil + } + do { + let audioLoadStart = Date() + let audio = try AudioConverter().resampleAudioFile(path: audioPath) + let audioLoadTime = Date().timeIntervalSince(audioLoadStart) + let duration = Float(audio.count) / 16000.0 + print(" \(meeting): \(String(format: "%.1f", duration))s") + + var mask: [Bool]? = nil + if let vadThreshold { + mask = try await speechMask(audio: audio, threshold: vadThreshold, vad: vad) + } + + let start = Date() + let (probs, frames) = try diarizer.processComplete(audio, speechMask: mask) + let processingTime = Date().timeIntervalSince(start) + let rtfx = duration / Float(processingTime) + if vadThreshold != nil { + let p = diarizer.lastProfile + print( + " VAD skipped \(p.skippedChunks)/\(p.chunkCount + p.skippedChunks) chunks") + } + + let segments = Nemotron3Diarizer.segments( + probabilities: probs, frameCount: frames, threshold: threshold, minDurationSeconds: 0) + + var groundTruth: [TimedSpeakerSegment] = [] + if let rttmURL = DiarizationBenchmarkUtils.getRTTMURL(for: meeting, dataset: dataset), + FileManager.default.fileExists(atPath: rttmURL.path), + let content = try? String(contentsOf: rttmURL, encoding: .utf8) + { + groundTruth = parseRTTM(content) + } + if groundTruth.isEmpty, dataset == .ami { + groundTruth = try AMIParser.loadWordAlignedGroundTruth(for: meeting, duration: duration) + } + guard !groundTruth.isEmpty else { + print(" No ground truth for \(meeting)") + return nil + } + + let ref = groundTruth.map { + DERSpeakerSegment( + speaker: $0.speakerId, start: Double($0.startTimeSeconds), end: Double($0.endTimeSeconds)) + } + let hyp = segments.map { + DERSpeakerSegment( + speaker: "speaker_\($0.speakerIndex)", start: Double($0.startSeconds), + end: Double($0.endSeconds)) + } + let der = DiarizationDER.compute(ref: ref, hyp: hyp, frameStep: 0.01, collar: collar) + let totalRef = max(der.totalRefSpeech, .leastNonzeroMagnitude) + + return DiarizationBenchmarkUtils.BenchmarkResult( + meetingName: meeting, + der: Float(der.der * 100), + missRate: Float(der.miss / totalRef * 100), + falseAlarmRate: Float(der.falseAlarm / totalRef * 100), + speakerErrorRate: Float(der.confusion / totalRef * 100), + rtfx: rtfx, + processingTime: processingTime, + totalFrames: frames, + detectedSpeakers: Set(segments.map(\.speakerIndex)).count, + groundTruthSpeakers: Set(groundTruth.map(\.speakerId)).count, + modelLoadTime: 0, + audioLoadTime: audioLoadTime + ) + } catch { + print(" Error on \(meeting): \(error)") + return nil + } + } + + private static func parseRTTM(_ content: String) -> [TimedSpeakerSegment] { + var segments: [TimedSpeakerSegment] = [] + for line in content.components(separatedBy: .newlines) { + let parts = line.trimmingCharacters(in: .whitespaces) + .components(separatedBy: .whitespaces).filter { !$0.isEmpty } + guard parts.count >= 8, parts[0] == "SPEAKER", + let start = Float(parts[3]), let dur = Float(parts[4]) + else { continue } + segments.append( + TimedSpeakerSegment( + speakerId: parts[7], embedding: [], startTimeSeconds: start, + endTimeSeconds: start + dur, qualityScore: 1.0)) + } + return segments + } +} + +extension Array { + fileprivate subscript(safe index: Int) -> Element? { + indices.contains(index) ? self[index] : nil + } +} +#endif diff --git a/Sources/FluidAudioCLI/Commands/SortformerBenchmark.swift b/Sources/FluidAudioCLI/Commands/SortformerBenchmark.swift index 26e7c46f1..ce52283fb 100644 --- a/Sources/FluidAudioCLI/Commands/SortformerBenchmark.swift +++ b/Sources/FluidAudioCLI/Commands/SortformerBenchmark.swift @@ -649,7 +649,7 @@ enum SortformerBenchmark { switch dataset { case .ami: groundTruthSpeakers = AMIParser.getGroundTruthSpeakerCount(for: meetingName) - case .voxconverse, .callhome: + default: // Count unique speakers from ground truth groundTruthSpeakers = Set(groundTruth.map { $0.speakerId }).count } diff --git a/Sources/FluidAudioCLI/FluidAudioCLI.swift b/Sources/FluidAudioCLI/FluidAudioCLI.swift index 057ba1263..c46426afa 100644 --- a/Sources/FluidAudioCLI/FluidAudioCLI.swift +++ b/Sources/FluidAudioCLI/FluidAudioCLI.swift @@ -68,6 +68,12 @@ struct FluidAudioCLI { await EmissionDelayBenchmark.runCLI(arguments: Array(arguments.dropFirst(2))) case "sortformer": await SortformerCommand.run(arguments: Array(arguments.dropFirst(2))) + case "nemotron3-diarize": + await Nemotron3DiarizeCommand.runDiarize(arguments: Array(arguments.dropFirst(2))) + case "nemotron3-benchmark": + await Nemotron3DiarizeCommand.runBenchmark(arguments: Array(arguments.dropFirst(2))) + case "nemotron3-batch": + await Nemotron3DiarizeCommand.runBatch(arguments: Array(arguments.dropFirst(2))) case "sortformer-benchmark": await SortformerBenchmark.run(arguments: Array(arguments.dropFirst(2))) case "lseend": diff --git a/Tests/FluidAudioTests/Diarizer/Nemotron3/Nemotron3StateUpdaterTests.swift b/Tests/FluidAudioTests/Diarizer/Nemotron3/Nemotron3StateUpdaterTests.swift new file mode 100644 index 000000000..7b2381747 --- /dev/null +++ b/Tests/FluidAudioTests/Diarizer/Nemotron3/Nemotron3StateUpdaterTests.swift @@ -0,0 +1,215 @@ +import Foundation +import XCTest + +@testable import FluidAudio + +final class Nemotron3StateUpdaterTests: XCTestCase { + + private var config: Nemotron3Config { .low } + + private func makeUpdater() -> Nemotron3StateUpdater { + Nemotron3StateUpdater( + config: config, + silenceEmbedding: [Float](repeating: 0.01, count: config.preEncoderDims)) + } + + /// Packed predictions sized for the current state: [spkcache | fifo | chunk] left-packed. + private func makePredictions( + state: Nemotron3StreamingState, chunkLen: Int, value: Float = 0.9 + ) -> (preds: [Float], hires: [Float]) { + let s = config.numSpeakers + let packed = config.packedFrames + var preds = [Float](repeating: 0, count: packed * s) + let valid = state.spkcacheLength + state.fifoLength + chunkLen + for frame in 0.. [Float] { + [Float](repeating: fill, count: config.chunkEncFrames * config.preEncoderDims) + } + + // MARK: - FIFO accumulation + + func testFirstChunkGoesToFifo() throws { + let updater = makeUpdater() + var state = Nemotron3StreamingState(config: config) + let chunkLen = config.chunkLen // full chunk, rc consumed + let (preds, hires) = makePredictions(state: state, chunkLen: chunkLen) + + let result = try updater.update( + state: &state, + chunkEmbeddings: makeChunkEmbeddings(frames: config.chunkEncFrames), + chunkEncLength: config.chunkEncFrames, + predictions: preds, + highResPredictions: hires, + lc: 0, + rc: config.chunkRightContext + ) + + XCTAssertEqual(state.fifoLength, chunkLen, "core frames should land in FIFO") + XCTAssertEqual(state.spkcacheLength, 0, "no cache update before FIFO overflow") + XCTAssertEqual(result.frameCount, chunkLen * config.upsampleFactor, "10ms output per core frame") + XCTAssertEqual(result.probabilities.count, result.frameCount * config.numSpeakers) + XCTAssertEqual(result.probabilities[0], 0.9, accuracy: 1e-6) + } + + func testFifoPopMovesFramesToSpkcache() throws { + let updater = makeUpdater() + var state = Nemotron3StreamingState(config: config) + + // Fill FIFO just below capacity, then push one more chunk to trigger a pop. + var steps = 0 + while state.fifoLength + config.chunkLen <= config.fifoLen { + let (preds, hires) = makePredictions(state: state, chunkLen: config.chunkLen) + _ = try updater.update( + state: &state, + chunkEmbeddings: makeChunkEmbeddings(frames: config.chunkEncFrames), + chunkEncLength: config.chunkEncFrames, + predictions: preds, highResPredictions: hires, + lc: 0, rc: config.chunkRightContext) + steps += 1 + } + XCTAssertEqual(state.spkcacheLength, 0) + let fifoBefore = state.fifoLength + + let (preds, hires) = makePredictions(state: state, chunkLen: config.chunkLen) + _ = try updater.update( + state: &state, + chunkEmbeddings: makeChunkEmbeddings(frames: config.chunkEncFrames), + chunkEncLength: config.chunkEncFrames, + predictions: preds, highResPredictions: hires, + lc: 0, rc: config.chunkRightContext) + + // NeMo pop rule: pop = min(combined, max(updatePeriod, overflow)) + let combined = fifoBefore + config.chunkLen + let expectedPop = min(combined, max(config.spkcacheUpdatePeriod, combined - config.fifoLen)) + XCTAssertEqual(state.spkcacheLength, expectedPop) + XCTAssertEqual(state.fifoLength, combined - expectedPop) + } + + func testZeroChunkFlushesFifo() throws { + let updater = makeUpdater() + var state = Nemotron3StreamingState(config: config) + + let (preds1, hires1) = makePredictions(state: state, chunkLen: config.chunkLen) + _ = try updater.update( + state: &state, + chunkEmbeddings: makeChunkEmbeddings(frames: config.chunkEncFrames), + chunkEncLength: config.chunkEncFrames, + predictions: preds1, highResPredictions: hires1, + lc: 0, rc: config.chunkRightContext) + let fifoBefore = state.fifoLength + XCTAssertGreaterThan(fifoBefore, 0) + + let (preds2, hires2) = makePredictions(state: state, chunkLen: 0) + let result = try updater.update( + state: &state, + chunkEmbeddings: makeChunkEmbeddings(frames: config.chunkEncFrames), + chunkEncLength: 0, + predictions: preds2, highResPredictions: hires2, + lc: 0, rc: 0) + + XCTAssertEqual(state.fifoLength, 0, "zero-length chunk must flush the FIFO") + XCTAssertEqual(state.spkcacheLength, fifoBefore, "flushed frames land in the cache") + XCTAssertEqual(result.frameCount, 0) + } + + // MARK: - Compression + + func testCompressionCapsSpkcacheAtCapacity() throws { + let updater = makeUpdater() + var state = Nemotron3StreamingState(config: config) + + // Stream enough active chunks to overflow the speaker cache. + // Each pop moves updatePeriod (222) frames; capacity 264 -> second pop compresses. + var iterations = 0 + while !state.spkcacheCompressed && iterations < 200 { + let (preds, hires) = makePredictions(state: state, chunkLen: config.chunkLen) + _ = try updater.update( + state: &state, + chunkEmbeddings: makeChunkEmbeddings(frames: config.chunkEncFrames), + chunkEncLength: config.chunkEncFrames, + predictions: preds, highResPredictions: hires, + lc: 0, rc: config.chunkRightContext) + iterations += 1 + XCTAssertLessThanOrEqual(state.spkcacheLength, config.spkcacheLen) + } + XCTAssertTrue(state.spkcacheCompressed, "cache should compress after sustained speech") + XCTAssertEqual(state.spkcacheLength, config.spkcacheLen) + } + + func testCompressionInsertsSilenceEmbeddingForDisabledSlots() throws { + let updater = makeUpdater() + var state = Nemotron3StreamingState(config: config) + + // All-silence predictions: every score disables, so compression fills slots with the + // learned silence embedding. + var iterations = 0 + while !state.spkcacheCompressed && iterations < 200 { + let (_, hires) = makePredictions(state: state, chunkLen: config.chunkLen, value: 0.0) + let silent = [Float](repeating: 0, count: config.packedFrames * config.numSpeakers) + _ = try updater.update( + state: &state, + chunkEmbeddings: makeChunkEmbeddings(frames: config.chunkEncFrames, fill: 0.7), + chunkEncLength: config.chunkEncFrames, + predictions: silent, highResPredictions: hires, + lc: 0, rc: config.chunkRightContext) + iterations += 1 + } + XCTAssertTrue(state.spkcacheCompressed) + // All frames were silent -> every selected slot should carry the silence embedding. + XCTAssertEqual(state.spkcache[0], 0.01, accuracy: 1e-6) + // Predictions for silence slots are zeroed. + XCTAssertEqual(state.spkcachePreds[0], 0, accuracy: 1e-6) + } + + // MARK: - Config invariants + + func testPresetShapes() { + XCTAssertEqual(Nemotron3Config.low.chunkMelFrames, 104) + XCTAssertEqual(Nemotron3Config.low.packedFrames, 541) + XCTAssertEqual(Nemotron3Config.veryLow.chunkMelFrames, 64) + XCTAssertEqual(Nemotron3Config.veryLow.packedFrames, 536) + XCTAssertEqual(Nemotron3Config.ultraLow.chunkMelFrames, 32) + XCTAssertEqual(Nemotron3Config.ultraLow.packedFrames, 532) + XCTAssertEqual(Nemotron3Config.offline.chunkMelFrames, 3040) + XCTAssertEqual(Nemotron3Config.offline.packedFrames, 684) + } + + func testPresetLookup() { + XCTAssertNotNil(Nemotron3Config.preset(named: "low")) + XCTAssertNotNil(Nemotron3Config.preset(named: "offline")) + XCTAssertNil(Nemotron3Config.preset(named: "bogus")) + } +} + +final class Nemotron3FeatureLoaderTests: XCTestCase { + + func testLoaderEmitsTailWithShrunkRightContext() { + let config = Nemotron3Config.low + let mel = config.melFeatures + // 2.5 core chunks of mel frames, no full right context at the tail. + let core = config.chunkLen * config.subsamplingFactor + let frames = core * 2 + core / 2 + let featSeq = [Float](repeating: 1, count: frames * mel) + + var loader = Nemotron3FeatureLoader( + config: config, featSeq: featSeq, featLength: frames, featSeqLength: frames) + var chunks: [(length: Int, right: Int)] = [] + while let c = loader.next() { + chunks.append((c.length, c.rightOffset)) + XCTAssertEqual(c.features.count, config.chunkMelFrames * mel, "fixed capacity padding") + } + + XCTAssertEqual(chunks.count, 3, "tail chunk must still be emitted") + XCTAssertEqual(chunks[0].right, config.chunkRightContext * config.subsamplingFactor) + XCTAssertLessThan(chunks[2].right, config.chunkRightContext * config.subsamplingFactor) + } +} diff --git a/Tests/FluidAudioTests/Diarizer/Nemotron3/Nemotron3StreamingFrontendTests.swift b/Tests/FluidAudioTests/Diarizer/Nemotron3/Nemotron3StreamingFrontendTests.swift new file mode 100644 index 000000000..14a4b8399 --- /dev/null +++ b/Tests/FluidAudioTests/Diarizer/Nemotron3/Nemotron3StreamingFrontendTests.swift @@ -0,0 +1,141 @@ +import XCTest + +@testable import FluidAudio + +/// The audio-in streaming front end must reproduce the batch path exactly: the same +/// mel frames as center-padded extraction and the same chunk sequence as +/// `Nemotron3FeatureLoader`, regardless of how the audio is batched into `append`. +final class Nemotron3StreamingFrontendTests: XCTestCase { + + /// Deterministic audio with speech-like structure (tones + noise), sized to not be a + /// multiple of the mel hop so edge handling is exercised. + private func makeAudio(seconds: Double) -> [Float] { + let count = Int(16000 * seconds) + 137 + srand48(11) + return (0.. [Chunk] { + let (feat, featLength, featSeqLength) = AudioMelSpectrogram().computeFlatTransposed(audio: audio) + var loader = Nemotron3FeatureLoader( + config: config, featSeq: feat, featLength: featLength, featSeqLength: featSeqLength) + var chunks: [Chunk] = [] + while let c = loader.next() { chunks.append(c) } + return chunks + } + + private func streamChunks(config: Nemotron3Config, audio: [Float], batchSize: Int) -> [Chunk] { + var frontend = Nemotron3StreamingFrontend(config: config) + var chunks: [Chunk] = [] + var fed = 0 + while fed < audio.count { + let end = min(fed + batchSize, audio.count) + frontend.append(Array(audio[fed.. Bool { + let shape = array.shape.map(\.intValue) + let strides = array.strides.map(\.intValue) + var expected = 1 + for dim in stride(from: shape.count - 1, through: 0, by: -1) { + if strides[dim] != expected { return false } + expected *= shape[dim] + } + return true + } + + func testAlignedArraysPadNonTileAlignedInnermostDims() throws { + // Documents the underlying behavior this bug class depends on. If this ever + // starts failing (helper made contiguous), the guards below become moot — fine. + let optimizer = ANEMemoryOptimizer() + let padded = try optimizer.createAlignedArray(shape: [1, 8, 1], dataType: .float32) + XCTAssertFalse( + isContiguous(padded), + "expected [1, 8, 1] aligned array to be stride-padded; update layout assumptions") + } + + func testFloatsFromStridedArrayHonorsStrides() throws { + let optimizer = ANEMemoryOptimizer() + for t in awkwardSizes { + let array = try optimizer.createAlignedArray( + shape: [1, NSNumber(value: t), 1], dataType: .float32) + // Write via logical (stride-aware) subscripting. + for i in 0..