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

Documentation/Diarization/Nemotron3.md

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

Sources/FluidAudio/Diarizer/Nemotron3/Nemotron3Diarizer.swift

@@ -0,0 +1,257 @@
+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")
+
+    /// 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)
+    }
+
+    public func reset() {
+        state = Nemotron3StreamingState(config: config)
+    }
+
+    /// 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..<hi].contains(true)
+                if !hasSpeech {
+                    let lcEnc = (chunk.leftOffset + sub / 2) / sub
+                    let rcEnc = (chunk.rightOffset + sub - 1) / sub
+                    let encLen = (chunk.length + sub - 1) / sub
+                    let chunkFrames = min(
+                        max(encLen - lcEnc, 0), config.chunkEncFrames - lcEnc - rcEnc)
+                    total.append(
+                        contentsOf: repeatElement(
+                            0, count: chunkFrames * config.upsampleFactor * config.numSpeakers))
+                    profile.skippedChunks += 1
+                    coreStart = coreEnd
+                    return true
+                }
+            }
+            coreStart = coreEnd
+
+            tStage = Date()
+            let out =
+                config.splitGraph
+                ? try models.runSplit(
+                    chunk: chunk.features, chunkLength: chunk.length, state: state, config: config)
+                : try models.run(
+                    chunk: chunk.features, chunkLength: chunk.length, state: state, config: config)
+            profile.inferenceSeconds += Date().timeIntervalSince(tStage)
+            profile.inputPrepSeconds += out.inputPrepSeconds
+            profile.predictSeconds += out.predictSeconds
+            profile.readbackSeconds += out.readbackSeconds
+
+            tStage = Date()
+            let result = try updater.update(
+                state: &state,
+                chunkEmbeddings: out.chunkEmbeddings,
+                chunkEncLength: out.chunkLength,
+                predictions: out.predictions,
+                highResPredictions: out.highResPredictions,
+                lc: (chunk.leftOffset + sub / 2) / sub,
+                rc: (chunk.rightOffset + sub - 1) / sub
+            )
+            profile.stateUpdateSeconds += Date().timeIntervalSince(tStage)
+
+            tStage = Date()
+            total.append(contentsOf: result.probabilities)
+            profile.outputAppendSeconds += Date().timeIntervalSince(tStage)
+            profile.chunkCount += 1
+            return true
+        }) {}
+
+        // NeMo trims to ceil(mel_frames / output_subsampling_factor); output factor is 1 (10 ms).
+        let outputFrames = min(featSeqLength, total.count / config.numSpeakers)
+        profile.totalSeconds = Date().timeIntervalSince(t0)
+        lastProfile = profile
+        return (Array(total[0..<(outputFrames * config.numSpeakers)]), outputFrames)
+    }
+
+    /// 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..<numSpeakers {
+            var start: Int? = nil
+            for frame in 0...frameCount {
+                let active = frame < frameCount && probabilities[frame * numSpeakers + spk] > 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..<endIdx])
+        // Zero-pad to the model's fixed mel capacity.
+        if features.count < capacityMel * melFeatures {
+            features.append(contentsOf: repeatElement(0, count: capacityMel * melFeatures - features.count))
+        }
+        let frames = endFeat + rightOffset - (startFeat - leftOffset)
+        let length = max(min(featSeqLength - startFeat + leftOffset, frames), 0)
+
+        startFeat = endFeat
+        return (features, length, leftOffset, rightOffset)
+    }
+}