From ebc175391242974e1b1117a841a03ed436128633 Mon Sep 17 00:00:00 2001 From: Alex-Wengg Date: Sun, 30 Aug 2026 16:27:20 -0400 Subject: [PATCH 1/9] feat(diarizer): Nemotron 3 Diarization support (8-speaker streaming Sortformer) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Swift runtime for NVIDIA's Nemotron 3 Diarization (early access): 8 speakers, arrival-order output, 10 ms frame resolution, streaming and offline profiles. - Nemotron3Diarizer / Nemotron3StateUpdater: port of NeMo streaming_update_async at batch 1 — fixed-capacity speaker-cache/FIFO state, score-based cache compression with the checkpoint's learned silence embedding, first-vs-later compression prediction freezing, NeMo-exact tail chunking, and 10 ms high-resolution output extracted before state mutation. Closed-loop output verified against the NeMo reference (99.995% frame agreement on a 120 s fixture). - Nemotron3Models: local-directory CoreML loading (no HF path until the model's public release), stride-aware vDSP_mmov output readback (the outputs are fp16 with padded rows; naive reads silently scramble or run 40x slower), per-chunk autoreleasepool (long ANE runs otherwise exhaust the IOSurface pool), and an optional split-graph mode that runs feature stacking + the 1024->512 projection host-side for a 100% ANE-resident pure-fp transformer graph. - Presets for the model-card profiles (offline/low/verylow/ultra) plus chunk-ladder profiles (fast/fast24/fast32/fast128/efficient) and int8/split selectors; optional VAD gating via a per-frame speech mask for sparse audio. - CLI: nemotron3-diarize, nemotron3-benchmark (AMI/VoxConverse harness with compute-unit routing and sweep flags), nemotron3-batch (concurrent GPU workers). - Tests: state updater semantics (FIFO pop/flush, compression, silence slots), feature loader tail handling, and padded-stride tensor layout regressions. Model weights are not distributed with this change; they load from a local directory and will move to HuggingFace auto-download after NVIDIA's public release. --- .../Nemotron3/Nemotron3Diarizer.swift | 257 +++++++ .../Diarizer/Nemotron3/Nemotron3Models.swift | 368 ++++++++++ .../Nemotron3/Nemotron3StateUpdater.swift | 371 ++++++++++ .../Diarizer/Nemotron3/Nemotron3Types.swift | 265 +++++++ .../Commands/Nemotron3DiarizeCommand.swift | 669 ++++++++++++++++++ Sources/FluidAudioCLI/FluidAudioCLI.swift | 6 + .../Nemotron3StateUpdaterTests.swift | 215 ++++++ .../Nemotron3TensorLayoutTests.swift | 90 +++ 8 files changed, 2241 insertions(+) create mode 100644 Sources/FluidAudio/Diarizer/Nemotron3/Nemotron3Diarizer.swift create mode 100644 Sources/FluidAudio/Diarizer/Nemotron3/Nemotron3Models.swift create mode 100644 Sources/FluidAudio/Diarizer/Nemotron3/Nemotron3StateUpdater.swift create mode 100644 Sources/FluidAudio/Diarizer/Nemotron3/Nemotron3Types.swift create mode 100644 Sources/FluidAudioCLI/Commands/Nemotron3DiarizeCommand.swift create mode 100644 Tests/FluidAudioTests/Diarizer/Nemotron3/Nemotron3StateUpdaterTests.swift create mode 100644 Tests/FluidAudioTests/Diarizer/Nemotron3/Nemotron3TensorLayoutTests.swift diff --git a/Sources/FluidAudio/Diarizer/Nemotron3/Nemotron3Diarizer.swift b/Sources/FluidAudio/Diarizer/Nemotron3/Nemotron3Diarizer.swift new file mode 100644 index 00000000..aeb49b26 --- /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) + let predictSeconds = Date().timeIntervalSince(tStage) + tStage = Date() + + guard let predsArray = output.featureValue(for: "speaker_preds")?.multiArrayValue, + let hiresArray = output.featureValue(for: "speaker_preds_10ms")?.multiArrayValue + else { + throw Nemotron3Error.inferenceFailed("Missing split model outputs") + } + 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) + let predictSeconds = Date().timeIntervalSince(tStage) + tStage = Date() + + guard let predsArray = output.featureValue(for: "speaker_preds")?.multiArrayValue, + let hiresArray = output.featureValue(for: "speaker_preds_10ms")?.multiArrayValue, + let embsArray = output.featureValue(for: "chunk_pre_encode_embs")?.multiArrayValue + else { + throw Nemotron3Error.inferenceFailed("Missing model outputs") + } + 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 00000000..47ad3ae9 --- /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 00000000..d664184e --- /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 057ba126..c46426af 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 00000000..7b238174 --- /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 00000000..03ae9dba --- /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.. Date: Sun, 30 Aug 2026 16:37:29 -0400 Subject: [PATCH 2/9] docs(diarizer): Nemotron 3 preset guide and implementation notes Preset chooser table (size / audio chunk / latency / pros / cons), quick start, CLI reference, and implementation notes. Accuracy and throughput figures are deferred to the model's public release per the early-access evaluation terms. --- Documentation/Diarization/Nemotron3.md | 97 ++++++++++++++++++++++++++ 1 file changed, 97 insertions(+) create mode 100644 Documentation/Diarization/Nemotron3.md diff --git a/Documentation/Diarization/Nemotron3.md b/Documentation/Diarization/Nemotron3.md new file mode 100644 index 00000000..28939f17 --- /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. From b19f818b8b4ad7bc4555d06f372d2062199a8cd4 Mon Sep 17 00:00:00 2001 From: Alex-Wengg Date: Sun, 30 Aug 2026 17:29:50 -0400 Subject: [PATCH 3/9] perf(diarizer/nemotron3): preallocated output backings MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Route both prediction paths through MLPredictionOptions.outputBackings with preallocated contiguous fp16 arrays. Latency is unchanged (the residual ~1ms predict-vs-benchtool gap lives inside CoreML's dispatch, not output allocation), but per-call output IOSurface allocation is gone — removing the root cause of pool exhaustion on long ANE runs (the per-chunk autoreleasepool remains as defense in depth). Output verified bit-identical on the parity fixture for the split path and DER-identical on AMI for the monolithic path. --- .../Diarizer/Nemotron3/Nemotron3Models.swift | 57 +++++++++++++++---- 1 file changed, 45 insertions(+), 12 deletions(-) diff --git a/Sources/FluidAudio/Diarizer/Nemotron3/Nemotron3Models.swift b/Sources/FluidAudio/Diarizer/Nemotron3/Nemotron3Models.swift index 5918e3af..e88f1b87 100644 --- a/Sources/FluidAudio/Diarizer/Nemotron3/Nemotron3Models.swift +++ b/Sources/FluidAudio/Diarizer/Nemotron3/Nemotron3Models.swift @@ -26,6 +26,14 @@ public struct Nemotron3Models { 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( @@ -71,6 +79,31 @@ public struct Nemotron3Models { 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. @@ -196,15 +229,12 @@ public struct Nemotron3Models { let inputPrepSeconds = Date().timeIntervalSince(tStage) tStage = Date() - let output = try model.prediction(from: inputs) + let output = try model.prediction(from: inputs, options: predictionOptions) let predictSeconds = Date().timeIntervalSince(tStage) tStage = Date() - guard let predsArray = output.featureValue(for: "speaker_preds")?.multiArrayValue, - let hiresArray = output.featureValue(for: "speaker_preds_10ms")?.multiArrayValue - else { - throw Nemotron3Error.inferenceFailed("Missing split model outputs") - } + 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), @@ -264,15 +294,18 @@ public struct Nemotron3Models { let inputPrepSeconds = Date().timeIntervalSince(tStage) tStage = Date() - let output = try model.prediction(from: inputs) + let output = try model.prediction(from: inputs, options: predictionOptions) let predictSeconds = Date().timeIntervalSince(tStage) tStage = Date() - guard let predsArray = output.featureValue(for: "speaker_preds")?.multiArrayValue, - let hiresArray = output.featureValue(for: "speaker_preds_10ms")?.multiArrayValue, - let embsArray = output.featureValue(for: "chunk_pre_encode_embs")?.multiArrayValue - else { - throw Nemotron3Error.inferenceFailed("Missing model outputs") + // 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) From e4d2780489ed3bccec141ca1c6af5d8266ff3c9f Mon Sep 17 00:00:00 2001 From: Alex-Wengg Date: Wed, 2 Sep 2026 16:14:43 -0400 Subject: [PATCH 4/9] feat(cli/nemotron3-benchmark): full card-matrix dataset support + SCA/MAE Extends the benchmark harness to cover every dataset condition and latency profile of NVIDIA's Nemotron 3 model-card matrix that is obtainable without an LDC license: - New --dataset values: ami-sdm (Array1-01), alimeeting-far (array ch0), alimeeting-near (equal-weight headset mix), notsofar-mhm, notsofar-sc. Audio materialized by two new reproducible scripts (ffmpeg): AliMeeting card audio, and NOTSOFAR1 eval mixes + word-timing RTTMs (0.2 s merge, mirroring the nttcslab forced-alignment convention). - SCA / MAE (the card's speaker-counting metrics) computed in the summary. - Summary header reports the actual dataset instead of hardcoded AMI SDM. Card-protocol references: nttcslab-sp/diar-forced-alignment for AMI (both conditions) and AliMeeting; collar 0, overlap included. The pipeline was validated by re-running previously measured rows to identical results before adding new conditions. Benchmark numbers are withheld until NVIDIA's public release per the preview eval license. --- Scripts/materialize_alimeeting_card_audio.py | 83 ++++++++++++ Scripts/materialize_notsofar_card_audio.py | 120 ++++++++++++++++++ .../Commands/DiarizationBenchmarkUtils.swift | 101 ++++++++++++++- .../Commands/Nemotron3DiarizeCommand.swift | 9 +- .../Commands/SortformerBenchmark.swift | 2 +- 5 files changed, 305 insertions(+), 10 deletions(-) create mode 100644 Scripts/materialize_alimeeting_card_audio.py create mode 100644 Scripts/materialize_notsofar_card_audio.py diff --git a/Scripts/materialize_alimeeting_card_audio.py b/Scripts/materialize_alimeeting_card_audio.py new file mode 100644 index 00000000..f91daed2 --- /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 00000000..4e789510 --- /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/FluidAudioCLI/Commands/DiarizationBenchmarkUtils.swift b/Sources/FluidAudioCLI/Commands/DiarizationBenchmarkUtils.swift index 7bab553b..2340c510 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 index d664184e..2952681a 100644 --- a/Sources/FluidAudioCLI/Commands/Nemotron3DiarizeCommand.swift +++ b/Sources/FluidAudioCLI/Commands/Nemotron3DiarizeCommand.swift @@ -26,7 +26,9 @@ enum Nemotron3DiarizeCommand { --output Write RTTM hypothesis nemotron3-benchmark options: - --dataset ami (default: ami) + --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) @@ -544,11 +546,14 @@ enum Nemotron3DiarizeCommand { let avgFA = results.map(\.falseAlarmRate).reduce(0, +) / Float(results.count) let avgConf = results.map(\.speakerErrorRate).reduce(0, +) / Float(results.count) let avgRTFx = results.map(\.rtfx).reduce(0, +) / Float(results.count) - print("\n=== Nemotron 3 Diarization (\(variantName)) — AMI SDM test (\(results.count) files) ===") + let (sca, mae) = DiarizationBenchmarkUtils.speakerCountMetrics(results: results) + print( + "\n=== Nemotron 3 Diarization (\(variantName)) — \(dataset.rawValue) (\(results.count) files) ===") print("Avg DER: \(String(format: "%.2f", avgDER))%") print( " miss \(String(format: "%.2f", avgMiss))% fa \(String(format: "%.2f", avgFA))% " + "conf \(String(format: "%.2f", avgConf))%") + print("SCA: \(String(format: "%.2f", sca))% MAE: \(String(format: "%.4f", mae))") print("Avg RTFx: \(String(format: "%.1f", avgRTFx))x") if let outputFile { diff --git a/Sources/FluidAudioCLI/Commands/SortformerBenchmark.swift b/Sources/FluidAudioCLI/Commands/SortformerBenchmark.swift index 26e7c46f..ce52283f 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 } From 72745fce91d2ad95bae9019c8c868af73114c043 Mon Sep 17 00:00:00 2001 From: Alex-Wengg Date: Sun, 20 Sep 2026 14:27:44 -0400 Subject: [PATCH 5/9] feat(diarizer/nemotron3): HF download and streaming audio API MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Nemotron 3 presets now load from FluidInference/nemotron-3-diarization-coreml via Nemotron3Models.loadFromHuggingFace (one bundle per preset under {monolithic,split}/ plus the root .bin assets; local-directory load unchanged). Repo.nemotron3Diarization added to ModelNames. Nemotron3Diarizer gains an audio-in streaming path — appendAudio / processBufferedAudio / finishStream — backed by Nemotron3StreamingFrontend: incremental mel extraction with full ±256-sample STFT context (prePadded mode, pre-emphasis history carried across calls) and the Nemotron3FeatureLoader chunk cadence. A chunk runs once its right context is buffered. Verified bit-exact against processComplete on real audio (NVIDIA's 8-voice demo clip, fast32 and low: max |Δ| = 0.0 over 9760 frames). The final flush clamps its buffer trim to the received sample count: for lengths with (N + 112) mod 160 < 15 the right-padded target frame count otherwise put the trim point past the buffer and trapped. Unit tests cover chunk sequence, tail trimming, release timing, reset, and the length residues that used to trap. CLI: --models is optional (HF download otherwise), nemotron3-diarize --streaming drives the live path in 100 ms pieces for parity checks. --- .../Nemotron3/Nemotron3Diarizer.swift | 189 ++++++++++++++++++ .../Diarizer/Nemotron3/Nemotron3Models.swift | 74 ++++++- .../Diarizer/Nemotron3/Nemotron3Types.swift | 13 ++ Sources/FluidAudio/ModelNames.swift | 19 ++ .../Shared/AudioMelSpectrogram.swift | 4 +- .../Commands/Nemotron3DiarizeCommand.swift | 79 +++++--- .../Nemotron3StreamingFrontendTests.swift | 141 +++++++++++++ 7 files changed, 486 insertions(+), 33 deletions(-) create mode 100644 Tests/FluidAudioTests/Diarizer/Nemotron3/Nemotron3StreamingFrontendTests.swift diff --git a/Sources/FluidAudio/Diarizer/Nemotron3/Nemotron3Diarizer.swift b/Sources/FluidAudio/Diarizer/Nemotron3/Nemotron3Diarizer.swift index aeb49b26..9d452a63 100644 --- a/Sources/FluidAudio/Diarizer/Nemotron3/Nemotron3Diarizer.swift +++ b/Sources/FluidAudio/Diarizer/Nemotron3/Nemotron3Diarizer.swift @@ -15,6 +15,12 @@ public final class Nemotron3Diarizer { 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 @@ -39,10 +45,14 @@ public final class Nemotron3Diarizer { 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 @@ -144,6 +154,58 @@ public final class Nemotron3Diarizer { return (Array(total[0..<(outputFrames * config.numSpeakers)]), outputFrames) } + // MARK: - Streaming (audio in) + + /// Buffer 16 kHz mono samples for the streaming path. Call `processBufferedAudio()` + /// afterwards to run every chunk the buffered audio completes. + /// + /// The streaming path is frame-exact with `processComplete` on the same audio: mel + /// frames are computed with their full STFT context and the chunk cadence mirrors + /// `Nemotron3FeatureLoader`. A chunk runs once `latencySeconds` of audio past its + /// start is available. + public func appendAudio(_ samples: [Float]) { + precondition(!streamFinished, "appendAudio after finishStream; call reset() first") + frontend.append(samples) + } + + /// Run every chunk the buffered audio completes and return their results in order. + /// Each result covers `chunkSeconds` of audio at 10 ms per frame. + public func processBufferedAudio() throws -> [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: @@ -255,3 +317,130 @@ public struct Nemotron3FeatureLoader { return (features, length, leftOffset, rightOffset) } } + +// MARK: - Streaming Frontend + +/// Audio-in front end for the streaming path: incremental mel extraction that is +/// frame-exact with center-padded batch extraction, plus the chunk cadence of +/// `Nemotron3FeatureLoader`. Model-free so the cadence and framing are unit-testable. +struct Nemotron3StreamingFrontend { + private let config: Nemotron3Config + private let mel = AudioMelSpectrogram() + + private var audio: [Float] = [] + /// Absolute sample index of `audio[0]`. + private var audioStart = 0 + private var samplesReceived = 0 + + private var melCache: [Float] = [] + /// Absolute mel-frame index of `melCache`'s first frame. + private var melCacheStart = 0 + private(set) var melFramesComputed = 0 + private var nextCoreMel = 0 + + init(config: Nemotron3Config) { + self.config = config + } + + mutating func reset() { + audio.removeAll(keepingCapacity: true) + audioStart = 0 + samplesReceived = 0 + melCache.removeAll(keepingCapacity: true) + melCacheStart = 0 + melFramesComputed = 0 + nextCoreMel = 0 + } + + mutating func append(_ samples: [Float]) { + audio.append(contentsOf: samples) + samplesReceived += samples.count + } + + /// The next chunk in `Nemotron3FeatureLoader` layout, or nil when the buffered audio + /// does not complete one. With `final`, the stream is right-padded like center-mode + /// extraction and the trailing (short) chunks are emitted. + mutating func nextChunk(final: Bool) -> (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. + 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 src = copyStart - audioStart + slice.withUnsafeMutableBufferPointer { dst in + audio.withUnsafeBufferPointer { buf in + dst.baseAddress!.advanced(by: copyStart - sliceStart) + .update(from: buf.baseAddress!.advanced(by: src), count: copyEnd - copyStart) + } + } + } + let previous = sliceStart - 1 + let lastSample: Float = previous >= 0 && previous < received ? audio[previous - audioStart] : 0 + + let (frames, _, _) = mel.computeFlatTransposed( + audio: slice, lastAudioSample: lastSample, 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 index e88f1b87..0669e1d7 100644 --- a/Sources/FluidAudio/Diarizer/Nemotron3/Nemotron3Models.swift +++ b/Sources/FluidAudio/Diarizer/Nemotron3/Nemotron3Models.swift @@ -106,22 +106,82 @@ public struct Nemotron3Models { self.predictionOptions = options } - /// Load from a local models directory. + /// 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 = directory.appendingPathComponent(config.modelFileName) + var modelURL = requestedModelURL 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")) + let packageURL = modelURL.deletingPathExtension().appendingPathExtension("mlpackage") guard FileManager.default.fileExists(atPath: packageURL.path) else { throw Nemotron3Error.modelLoadFailed( - "Neither \(config.modelFileName) nor its .mlpackage found in \(directory.path)") + "Neither \(config.modelFileName) nor its .mlpackage found at \(modelURL.deletingLastPathComponent().path)" + ) } modelURL = try await MLModel.compileModel(at: packageURL) } @@ -130,7 +190,7 @@ public struct Nemotron3Models { mlConfig.computeUnits = computeUnits let model = try MLModel(contentsOf: modelURL, configuration: mlConfig) - let silURL = directory.appendingPathComponent("learnable_sil_emb.bin") + 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)") } @@ -143,7 +203,7 @@ public struct Nemotron3Models { var projection: [Float]? = nil if config.splitGraph { - let projURL = directory.appendingPathComponent("pre_encode_proj_t.bin") + let projURL = directory.appendingPathComponent(ModelNames.Nemotron3.preEncodeProjectionFile) guard let projData = try? Data(contentsOf: projURL), projData.count == 1024 * 512 * MemoryLayout.size else { diff --git a/Sources/FluidAudio/Diarizer/Nemotron3/Nemotron3Types.swift b/Sources/FluidAudio/Diarizer/Nemotron3/Nemotron3Types.swift index 47ad3ae9..f057bcaa 100644 --- a/Sources/FluidAudio/Diarizer/Nemotron3/Nemotron3Types.swift +++ b/Sources/FluidAudio/Diarizer/Nemotron3/Nemotron3Types.swift @@ -74,6 +74,19 @@ public struct Nemotron3Config: Sendable { /// 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. diff --git a/Sources/FluidAudio/ModelNames.swift b/Sources/FluidAudio/ModelNames.swift index e2085d59..d8138bb1 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 dd629bc4..f46e403f 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 diff --git a/Sources/FluidAudioCLI/Commands/Nemotron3DiarizeCommand.swift b/Sources/FluidAudioCLI/Commands/Nemotron3DiarizeCommand.swift index 2952681a..2109e03b 100644 --- a/Sources/FluidAudioCLI/Commands/Nemotron3DiarizeCommand.swift +++ b/Sources/FluidAudioCLI/Commands/Nemotron3DiarizeCommand.swift @@ -3,25 +3,29 @@ import CoreML import FluidAudio import Foundation -/// CLI for Nemotron 3 Diarization preview (local eval-license models — no HF download). +/// CLI for Nemotron 3 Diarization. enum Nemotron3DiarizeCommand { private static let logger = AppLogger(category: "Nemotron3CLI") static func printUsage() { let usage = """ - Nemotron 3 Diarization (preview, internal evaluation only) + Nemotron 3 Diarization (8-speaker streaming) Usage: - fluidaudiocli nemotron3-diarize --models [options] - fluidaudiocli nemotron3-benchmark --models [options] + fluidaudiocli nemotron3-diarize [options] + fluidaudiocli nemotron3-benchmark [options] Shared options: - --models Directory containing Nemotron3Diarizer_.mlmodelc - and learnable_sil_emb.bin (REQUIRED) - --variant offline | low | verylow | ultra (default: low) + --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 @@ -47,7 +51,7 @@ enum Nemotron3DiarizeCommand { } private static func loadDiarizer( - modelsDir: String, variantName: String, custom: CustomShape = CustomShape(), + modelsDir: String?, variantName: String, custom: CustomShape = CustomShape(), computeUnits: MLComputeUnits = .all ) async throws -> (Nemotron3Diarizer, TimeInterval) { var config: Nemotron3Config @@ -64,14 +68,44 @@ enum Nemotron3DiarizeCommand { modelFileName: "Nemotron3Diarizer_\(variantName).mlmodelc") } let start = Date() - let models = try await Nemotron3Models.load( - config: config, - directory: URL(fileURLWithPath: modelsDir), - computeUnits: computeUnits - ) + 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 @@ -144,6 +178,7 @@ enum Nemotron3DiarizeCommand { var outputPath: String? var showProfile = false var useVad = false + var useStreaming = false var vadThreshold: Float = 0.85 var custom = CustomShape() var computeUnits: MLComputeUnits = .all @@ -169,6 +204,8 @@ enum Nemotron3DiarizeCommand { outputPath = arguments[safe: i] case "--profile": showProfile = true + case "--streaming": + useStreaming = true case "--vad": useVad = true case "--vad-threshold": @@ -201,7 +238,7 @@ enum Nemotron3DiarizeCommand { i += 1 } - guard let audioPath, let modelsDir else { + guard let audioPath else { printUsage() exit(1) } @@ -225,7 +262,10 @@ enum Nemotron3DiarizeCommand { } let start = Date() - let (probs, frames) = try diarizer.processComplete(audio, speechMask: mask) + let (probs, frames) = + useStreaming + ? try streamAudio(audio, through: diarizer) + : try diarizer.processComplete(audio, speechMask: mask) let elapsed = Date().timeIntervalSince(start) let rtfx = duration / Float(elapsed) if useVad { @@ -349,10 +389,6 @@ enum Nemotron3DiarizeCommand { i += 1 } - guard let modelsDir else { - printUsage() - exit(1) - } if files.isEmpty { files = DiarizationBenchmarkUtils.getFiles(for: dataset, maxFiles: maxFiles) } @@ -498,11 +534,6 @@ enum Nemotron3DiarizeCommand { i += 1 } - guard let modelsDir else { - printUsage() - exit(1) - } - do { let (diarizer, loadTime) = try await loadDiarizer( modelsDir: modelsDir, variantName: variantName, custom: custom, computeUnits: computeUnits) diff --git a/Tests/FluidAudioTests/Diarizer/Nemotron3/Nemotron3StreamingFrontendTests.swift b/Tests/FluidAudioTests/Diarizer/Nemotron3/Nemotron3StreamingFrontendTests.swift new file mode 100644 index 00000000..14a4b839 --- /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.. Date: Sun, 20 Sep 2026 14:27:44 -0400 Subject: [PATCH 6/9] feat(asr/sensevoice): transcribeDetailed keeps the model's language/emotion/event tags SenseVoiceManager.transcribeDetailed returns the text plus the leading <|...|> query tags (detected language among zh/en/yue/ja/ko/nospeech) so callers can route by language; transcribe(audio:) is unchanged. --- .../ASR/SenseVoice/SenseVoiceManager.swift | 45 ++++++++++++++++--- 1 file changed, 39 insertions(+), 6 deletions(-) diff --git a/Sources/FluidAudio/ASR/SenseVoice/SenseVoiceManager.swift b/Sources/FluidAudio/ASR/SenseVoice/SenseVoiceManager.swift index 6fcfb47b..4203e8d2 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] +} From 6f36a0a75d3f6fc3aa079b2d33ba0270864f9173 Mon Sep 17 00:00:00 2001 From: Alex-Wengg Date: Sun, 20 Sep 2026 14:27:44 -0400 Subject: [PATCH 7/9] docs(nemotron3): HF download quick start and streaming section --- Documentation/Diarization/Nemotron3.md | 59 +++++++++++++++++++------- 1 file changed, 43 insertions(+), 16 deletions(-) diff --git a/Documentation/Diarization/Nemotron3.md b/Documentation/Diarization/Nemotron3.md index 28939f17..928dc0c3 100644 --- a/Documentation/Diarization/Nemotron3.md +++ b/Documentation/Diarization/Nemotron3.md @@ -4,11 +4,11 @@ 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. +> **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 @@ -16,10 +16,7 @@ resolution, streaming and offline profiles from a single checkpoint. import FluidAudio let config = Nemotron3Config.fast32 // recommended default -let models = try await Nemotron3Models.load( - config: config, - directory: localModelsDirectoryURL -) +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 @@ -34,6 +31,31 @@ preserving the output timeline): 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 @@ -67,19 +89,24 @@ 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 +# 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 --models --variant fast32 --collar 0 +swift run fluidaudiocli nemotron3-benchmark --variant fast32 --collar 0 # Batch processing with concurrent GPU workers -swift run fluidaudiocli nemotron3-batch --models --workers 2 --files a,b,c +swift run fluidaudiocli nemotron3-batch --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. +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 From 50f4b98a22e8910d476876efdf1e3bbbd577817c Mon Sep 17 00:00:00 2001 From: Alex-Wengg Date: Sun, 20 Sep 2026 10:46:43 -0400 Subject: [PATCH 8/9] =?UTF-8?q?docs(ane-profiler):=20streaming=20diarizati?= =?UTF-8?q?on=20table=20=E2=80=94=20Sortformer=20v2.1=20vs=20Nemotron=203?= =?UTF-8?q?=20device=20split=20+=20per-call=20latency?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit (cherry picked from commit 53104aa506d6235a1fb12f07565a7779906e8c6b) --- Documentation/ANE_Profiler.md | 45 +++++++++++++++++++++++++++++++++++ 1 file changed, 45 insertions(+) diff --git a/Documentation/ANE_Profiler.md b/Documentation/ANE_Profiler.md index 9f536e98..65c75291 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 | From eec575e0df47952ed2eeb7a76009be56553a5d63 Mon Sep 17 00:00:00 2001 From: Alex-Wengg Date: Sun, 20 Sep 2026 15:16:26 -0400 Subject: [PATCH 9/9] fix(nemotron3): pre-emphasize before padding in the streaming final flush The batch mel path applies pre-emphasis to the audio and then zero-pads both edges, so the padding stays zero. The streaming front end built a zero-padded slice first and let the extractor pre-emphasize it, which left a -0.97 * x[N-1] spike at the first padded sample of the final flush. Every frame whose window covers that sample (the last one or two frames of a stream) diverged from the batch loader by up to ~6 in log-mel. Pre-emphasize the received samples in the front end (vDSP_vsma over the same values as the batch path, seeded from the retained history sample) and hand the extractor a preemph-free slice. Streaming chunks are now bit-identical to the batch loader for every chunk, including the tail. AudioMelSpectrogram gains a public `defaultPreemph` constant so the front end and the extractor share the coefficient. Caught by Nemotron3StreamingFrontendTests on CI (three parity tests failed on the last chunk); XCTest is unavailable locally so the tests only ran there. --- .../Nemotron3/Nemotron3Diarizer.swift | 30 ++++++++++++++----- .../Shared/AudioMelSpectrogram.swift | 5 +++- 2 files changed, 26 insertions(+), 9 deletions(-) diff --git a/Sources/FluidAudio/Diarizer/Nemotron3/Nemotron3Diarizer.swift b/Sources/FluidAudio/Diarizer/Nemotron3/Nemotron3Diarizer.swift index 9d452a63..ce1498a1 100644 --- a/Sources/FluidAudio/Diarizer/Nemotron3/Nemotron3Diarizer.swift +++ b/Sources/FluidAudio/Diarizer/Nemotron3/Nemotron3Diarizer.swift @@ -1,3 +1,4 @@ +import Accelerate import Foundation /// Streaming 8-speaker diarizer backed by NVIDIA's Nemotron 3 Diarization preview. @@ -325,7 +326,9 @@ public struct Nemotron3FeatureLoader { /// `Nemotron3FeatureLoader`. Model-free so the cadence and framing are unit-testable. struct Nemotron3StreamingFrontend { private let config: Nemotron3Config - private let mel = AudioMelSpectrogram() + /// Pre-emphasis is applied here, before padding, so the extractor runs without it. + private let mel = AudioMelSpectrogram(preemph: 0) + private let preemph = AudioMelSpectrogram.defaultPreemph private var audio: [Float] = [] /// Absolute sample index of `audio[0]`. @@ -410,27 +413,38 @@ struct Nemotron3StreamingFrontend { let done = melFramesComputed guard target > done else { return } - // Frame f is centered on sample f*hop; its window spans ±half around it. + // 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 - dst.baseAddress!.advanced(by: copyStart - sliceStart) - .update(from: buf.baseAddress!.advanced(by: src), count: copyEnd - copyStart) + 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 previous = sliceStart - 1 - let lastSample: Float = previous >= 0 && previous < received ? audio[previous - audioStart] : 0 let (frames, _, _) = mel.computeFlatTransposed( - audio: slice, lastAudioSample: lastSample, paddingMode: .prePadded, - expectedFrameCount: target - done) + audio: slice, paddingMode: .prePadded, expectedFrameCount: target - done) melCache.append(contentsOf: frames[0..<((target - done) * config.melFeatures)]) melFramesComputed = target diff --git a/Sources/FluidAudio/Shared/AudioMelSpectrogram.swift b/Sources/FluidAudio/Shared/AudioMelSpectrogram.swift index f46e403f..1b6ce956 100644 --- a/Sources/FluidAudio/Shared/AudioMelSpectrogram.swift +++ b/Sources/FluidAudio/Shared/AudioMelSpectrogram.swift @@ -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,