diff --git a/Documentation/Diarization/Nemotron3.md b/Documentation/Diarization/Nemotron3.md new file mode 100644 index 000000000..28939f17d --- /dev/null +++ b/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 --variant fast32 + +# Benchmark against AMI / VoxConverse harnesses +swift run fluidaudiocli nemotron3-benchmark --models --variant fast32 --collar 0 + +# Batch processing with concurrent GPU workers +swift run fluidaudiocli nemotron3-batch --models --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. diff --git a/Sources/FluidAudio/Diarizer/Nemotron3/Nemotron3Diarizer.swift b/Sources/FluidAudio/Diarizer/Nemotron3/Nemotron3Diarizer.swift new file mode 100644 index 000000000..aeb49b263 --- /dev/null +++ b/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.. 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.. Nemotron3Models { + let start = Date() + + var modelURL = directory.appendingPathComponent(config.modelFileName) + if !FileManager.default.fileExists(atPath: modelURL.path) { + // Fall back to the uncompiled mlpackage next to the expected mlmodelc. + let packageURL = directory.appendingPathComponent( + config.modelFileName.replacingOccurrences(of: ".mlmodelc", with: ".mlpackage")) + guard FileManager.default.fileExists(atPath: packageURL.path) else { + throw Nemotron3Error.modelLoadFailed( + "Neither \(config.modelFileName) nor its .mlpackage found in \(directory.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("learnable_sil_emb.bin") + 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("pre_encode_proj_t.bin") + 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..47ad3ae96 --- /dev/null +++ b/Sources/FluidAudio/Diarizer/Nemotron3/Nemotron3Types.swift @@ -0,0 +1,265 @@ +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 } + + // 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/FluidAudioCLI/Commands/Nemotron3DiarizeCommand.swift b/Sources/FluidAudioCLI/Commands/Nemotron3DiarizeCommand.swift new file mode 100644 index 000000000..d664184e7 --- /dev/null +++ b/Sources/FluidAudioCLI/Commands/Nemotron3DiarizeCommand.swift @@ -0,0 +1,669 @@ +#if os(macOS) +import CoreML +import FluidAudio +import Foundation + +/// CLI for Nemotron 3 Diarization preview (local eval-license models — no HF download). +enum Nemotron3DiarizeCommand { + private static let logger = AppLogger(category: "Nemotron3CLI") + + static func printUsage() { + let usage = """ + Nemotron 3 Diarization (preview, internal evaluation only) + + Usage: + fluidaudiocli nemotron3-diarize --models [options] + fluidaudiocli nemotron3-benchmark --models [options] + + Shared options: + --models Directory containing Nemotron3Diarizer_.mlmodelc + and learnable_sil_emb.bin (REQUIRED) + --variant offline | low | verylow | ultra (default: low) + --threshold Speaker activity threshold (default: 0.5) + + nemotron3-diarize options: + --dump-preds Write raw frame probabilities (float32 LE, [T, 8]) for parity checks + --output Write RTTM hypothesis + + nemotron3-benchmark options: + --dataset ami (default: ami) + --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 = try await Nemotron3Models.load( + config: config, + directory: URL(fileURLWithPath: modelsDir), + computeUnits: computeUnits + ) + return (Nemotron3Diarizer(config: config, models: models), Date().timeIntervalSince(start)) + } + + static func parseComputeUnits(_ s: String?) -> 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 + } + + guard let modelsDir else { + printUsage() + exit(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/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/Nemotron3TensorLayoutTests.swift b/Tests/FluidAudioTests/Diarizer/Nemotron3/Nemotron3TensorLayoutTests.swift new file mode 100644 index 000000000..03ae9dba8 --- /dev/null +++ b/Tests/FluidAudioTests/Diarizer/Nemotron3/Nemotron3TensorLayoutTests.swift @@ -0,0 +1,90 @@ +import CoreML +import Foundation +import XCTest + +@testable import FluidAudio + +/// Regression tests for the padded-stride MLMultiArray bug class. +/// +/// `ANEMemoryUtils.calculateOptimalStrides` pads innermost dimensions to tile +/// boundaries, so shapes like [1, T, 1] get a row stride of 16 — linear writes +/// through `dataPointer` then land at the wrong logical positions. This silently +/// corrupted the split-graph `output_mask` input (surfaced as ~90% frame agreement +/// instead of ~100%). These tests pin both directions: +/// - reads: `Nemotron3Models.floats(from:)` must honor strides for padded layouts, +/// - writes: buffers written linearly must actually be contiguous. +final class Nemotron3TensorLayoutTests: XCTestCase { + + /// Non-tile-aligned innermost sizes that trigger stride padding. + private let awkwardSizes = [1, 3, 5, 7, 9, 15, 17, 33, 340, 341] + + private func isContiguous(_ array: MLMultiArray) -> 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..