fix(diarizer/sortformer): consume BNNS-fixed v3 models + config-mismatch guard (#726)

The root-level Sortformer CoreML models hit a BNNS graph-compile crash on newer
BNNS ("tensor chunk_pre_encoder_embs_out as both an input and output"). The fixed
rebuild lives at v3/fp16/ in the HF repo; point ModelNames there so downloads pick
up the working models.

- ModelNames.Sortformer.modelsSubdirectory = "v3/fp16" (BNNS-fixed set); v3/palettized
  is the 6-bit, ~2.5x-smaller set for RAM-constrained devices.
- Add efficientV2_1 variant (chunk_len=25, ~2s latency, ~4x RTFx of fast) + config preset.
- SortformerDiarizer now validates the diarizer config against the model's embedded
  metadata on init and logs a clear error on mismatch (a mismatch silently produced
  incorrect/slow diarization — #726). spkcacheUpdatePeriod excluded (host-clamped).
- CLI: `sortformer --config fast|efficient|low|high`; `sortformer-benchmark --collar`,
  `--onset`, `--offset` (the hardcoded collar=0 / onset=0.5 skewed reported DER).

Author: Alex-Wengg

Sources/FluidAudio/Diarizer/Sortformer/SortformerDiarizer.swift

@@ -108,6 +108,8 @@ public final class SortformerDiarizer: Diarizer {
             mainModelPath: mainModelPath
         )
 
+        validateConfigMatch(loadedModels)
+
         // Use withLock helper to avoid direct NSLock usage in async context
         withLock {
             self._models = loadedModels
@@ -117,6 +119,34 @@ public final class SortformerDiarizer: Diarizer {
         logger.info("Sortformer initialized in \(String(format: "%.2f", loadedModels.compilationDuration))s")
     }
 
+    /// Warn loudly if the diarizer's `config` does not match the streaming parameters baked
+    /// into the loaded model. A mismatch (e.g. a `.default` config against a `highContextV2_1`
+    /// model) runs but produces incorrect and much slower results — issue #726.
+    private func validateConfigMatch(_ models: SortformerModels) {
+        guard let embedded = models.embeddedConfig else { return }
+        let current = SortformerModels.EmbeddedConfig(
+            chunkLen: config.chunkLen,
+            chunkLeftContext: config.chunkLeftContext,
+            chunkRightContext: config.chunkRightContext,
+            fifoLen: config.fifoLen,
+            spkcacheLen: config.spkcacheLen
+        )
+        guard current != embedded else { return }
+        logger.error(
+            """
+            Sortformer config mismatch — diarizer config does not match the loaded model. \
+            This produces incorrect and much slower diarization (issue #726). \
+            diarizer(chunkLen=\(current.chunkLen), leftCtx=\(current.chunkLeftContext), \
+            rightCtx=\(current.chunkRightContext), fifoLen=\(current.fifoLen), \
+            spkcacheLen=\(current.spkcacheLen)) \
+            vs model(chunkLen=\(embedded.chunkLen), leftCtx=\(embedded.chunkLeftContext), \
+            rightCtx=\(embedded.chunkRightContext), fifoLen=\(embedded.fifoLen), \
+            spkcacheLen=\(embedded.spkcacheLen)). \
+            Construct SortformerDiarizer with the SortformerConfig matching the model variant.
+            """
+        )
+    }
+
     /// Execute a closure while holding the lock
     private func withLock<T>(_ body: () throws -> T) rethrows -> T {
         lock.lock()
@@ -126,6 +156,8 @@ public final class SortformerDiarizer: Diarizer {
 
     /// Initialize with pre-loaded models.
     public func initialize(models: SortformerModels) {
+        validateConfigMatch(models)
+
         lock.lock()
         defer { lock.unlock() }
 

Sources/FluidAudio/Diarizer/Sortformer/SortformerModelInference.swift

@@ -156,6 +156,52 @@ extension SortformerModels {
     }
 }
 
+// MARK: - Embedded Configuration
+
+extension SortformerModels {
+
+    /// The model-shape-defining streaming parameters the converter embeds in the CoreML model
+    /// metadata. These determine the input tensor shapes and must match the host config.
+    ///
+    /// Note: `spkcache_update_period` is intentionally excluded — `SortformerConfig.init` clamps
+    /// it host-side (`max(min(period, fifoLen+chunkLen), chunkLen)`), so the host value legitimately
+    /// differs from the raw value baked into the model and is not a compatibility signal.
+    public struct EmbeddedConfig: Equatable, Sendable {
+        public let chunkLen: Int
+        public let chunkLeftContext: Int
+        public let chunkRightContext: Int
+        public let fifoLen: Int
+        public let spkcacheLen: Int
+    }
+
+    /// The variant-defining streaming parameters the converter writes into the CoreML model
+    /// metadata. Returns `nil` for older exports that don't carry them. Used to detect a
+    /// `SortformerConfig` that doesn't match the model — a mismatch yields incorrect and much
+    /// slower diarization (issue #726).
+    public var embeddedConfig: EmbeddedConfig? {
+        guard let meta = mainModel.modelDescription.metadata[.creatorDefinedKey] as? [String: String] else {
+            return nil
+        }
+        func value(_ key: String) -> Int? { meta[key].flatMap(Int.init) }
+        guard
+            let chunkLen = value("chunk_len"),
+            let chunkLeftContext = value("chunk_left_context"),
+            let chunkRightContext = value("chunk_right_context"),
+            let fifoLen = value("fifo_len"),
+            let spkcacheLen = value("spkcache_len")
+        else {
+            return nil
+        }
+        return EmbeddedConfig(
+            chunkLen: chunkLen,
+            chunkLeftContext: chunkLeftContext,
+            chunkRightContext: chunkRightContext,
+            fifoLen: fifoLen,
+            spkcacheLen: spkcacheLen
+        )
+    }
+}
+
 // MARK: - Main Model Inference
 
 extension SortformerModels {

Sources/FluidAudio/Diarizer/Sortformer/SortformerTypes.swift

@@ -194,6 +194,21 @@ public struct SortformerConfig: Sendable {
         spkcacheUpdatePeriod: 300
     )
 
+    /// Higher-throughput streaming config with Sortformer v2.1 weights (~2s output latency).
+    /// Same context as `fastV2_1` but a larger 25-frame chunk: per-inference cost is dominated
+    /// by the static speaker-cache + FIFO context, so a bigger chunk advances ~4x more audio per
+    /// call at near-identical latency (~4x real-time factor vs `fastV2_1`). Use when ~2s latency
+    /// is acceptable and throughput matters.
+    public static let efficientV2_1 = SortformerConfig(
+        modelVariant: .efficientV2_1,
+        chunkLen: 25,
+        chunkLeftContext: 1,
+        chunkRightContext: 7,
+        fifoLen: 40,
+        spkcacheLen: 188,
+        spkcacheUpdatePeriod: 31
+    )
+
     /// - Warning: If you don't use one of the default configurations, you must use a local model converted with that configuration.
     public init(
         modelVariant: ModelVariant? = .fastV2_1,

Sources/FluidAudio/ModelNames.swift

@@ -642,6 +642,9 @@ public enum ModelNames {
             case balancedV2_1
             case highContextV2
             case highContextV2_1
+            /// Higher-throughput streaming: larger chunk (~2s output latency) for ~4x the
+            /// real-time factor of `fastV2_1` at near-identical per-inference cost.
+            case efficientV2_1
 
             public var name: String {
                 switch self {
@@ -657,6 +660,8 @@ public enum ModelNames {
                     return "SortformerNvidiaHigh_v2"
                 case .highContextV2_1:
                     return "SortformerNvidiaHigh_v2.1"
+                case .efficientV2_1:
+                    return "SortformerEfficient_v2.1"
                 }
             }
 
@@ -674,18 +679,26 @@ public enum ModelNames {
                     return .highContextV2
                 case .highContextV2_1:
                     return .highContextV2_1
+                case .efficientV2_1:
+                    return .efficientV2_1
                 }
             }
 
             public var fileName: String {
-                return "\(name).mlmodelc"
+                return "\(Sortformer.modelsSubdirectory)/\(name).mlmodelc"
             }
 
             public func isCompatible(with config: SortformerConfig) -> Bool {
                 defaultConfiguration.isCompatible(with: config)
             }
         }
 
+        /// Repo subdirectory holding the active model set. `v3/fp16` is the BNNS-fixed rebuild
+        /// (the older root-level models hit a "tensor as both input and output" graph-compile
+        /// crash on newer BNNS — issue #726). Use `v3/palettized` for the 6-bit, ~2.5x-smaller
+        /// set (fixes RAM-driven crashes on older devices at ~+0.9pp DER).
+        public static let modelsSubdirectory = "v3/fp16"
+
         /// Lowest latency for streaming
         public static let defaultVariant: Variant = .fastV2_1
 

Sources/FluidAudioCLI/Commands/SortformerBenchmark.swift

@@ -68,6 +68,9 @@ enum SortformerBenchmark {
         var singleFile: String?
         var maxFiles: Int?
         var threshold: Float = 0.5
+        var collarSeconds: Double = 0
+        var onsetThreshold: Float?
+        var offsetThreshold: Float?
         var modelPath: String?
         var outputFile: String?
         var verbose = false
@@ -116,6 +119,21 @@ enum SortformerBenchmark {
                     threshold = Float(arguments[i + 1]) ?? 0.5
                     i += 1
                 }
+            case "--collar":
+                if i + 1 < arguments.count {
+                    collarSeconds = Double(arguments[i + 1]) ?? 0
+                    i += 1
+                }
+            case "--onset":
+                if i + 1 < arguments.count {
+                    onsetThreshold = Float(arguments[i + 1])
+                    i += 1
+                }
+            case "--offset":
+                if i + 1 < arguments.count {
+                    offsetThreshold = Float(arguments[i + 1])
+                    i += 1
+                }
             case "--model":
                 if i + 1 < arguments.count {
                     modelPath = arguments[i + 1]
@@ -294,7 +312,19 @@ enum SortformerBenchmark {
         if let v = weakBoostRate { config.weakBoostRate = v }
         if let v = minPosScoresRate { config.minPosScoresRate = v }
         if let v = spkcacheSilFramesPerSpk { config.spkcacheSilFramesPerSpk = v }
-        let diarizer = SortformerDiarizer(config: config)
+        // Allow overriding the timeline binarization thresholds (sortformerDefault = 0.5/0.5).
+        let diarizer: SortformerDiarizer
+        if onsetThreshold != nil || offsetThreshold != nil {
+            let timeline = DiarizerTimelineConfig(
+                numSpeakers: config.numSpeakers,
+                frameDurationSeconds: Float(config.frameDurationSeconds),
+                onsetThreshold: onsetThreshold ?? 0.5,
+                offsetThreshold: offsetThreshold ?? onsetThreshold ?? 0.5
+            )
+            diarizer = SortformerDiarizer(config: config, timelineConfig: timeline)
+        } else {
+            diarizer = SortformerDiarizer(config: config)
+        }
 
         do {
             if useHuggingFace {
@@ -340,6 +370,7 @@ enum SortformerBenchmark {
                 diarizer: diarizer,
                 modelLoadTime: modelLoadTime,
                 threshold: threshold,
+                collarSeconds: collarSeconds,
                 verbose: verbose
             )
 
@@ -381,6 +412,7 @@ enum SortformerBenchmark {
         diarizer: SortformerDiarizer,
         modelLoadTime: Double,
         threshold: Float,
+        collarSeconds: Double,
         verbose: Bool
     ) async -> BenchmarkResult? {
 
@@ -484,7 +516,7 @@ enum SortformerBenchmark {
                 ref: referenceSegments,
                 hyp: hypothesisSegments,
                 frameStep: derFrameStepSeconds,
-                collar: 0
+                collar: collarSeconds
             )
             let totalRefSpeech = max(derResult.totalRefSpeech, .leastNonzeroMagnitude)
             let derPercent = Float(derResult.der * 100)

Sources/FluidAudioCLI/Commands/SortformerCommand.swift

@@ -41,6 +41,7 @@ enum SortformerCommand {
         var weakBoostRate: Float?
         var minPosScoresRate: Float?
         var spkcacheSilFramesPerSpk: Int?
+        var configName = "default"
 
         // Parse remaining arguments
         var i = 1
@@ -88,6 +89,11 @@ enum SortformerCommand {
                     modelPath = arguments[i + 1]
                     i += 1
                 }
+            case "--config":
+                if i + 1 < arguments.count {
+                    configName = arguments[i + 1].lowercased()
+                    i += 1
+                }
             case "--threshold":
                 if i + 1 < arguments.count, let v = Float(arguments[i + 1]) {
                     predScoreThreshold = v
@@ -132,8 +138,20 @@ enum SortformerCommand {
         print("Sortformer Streaming Diarization")
         print("   Audio: \(audioFile)")
 
-        // Initialize Sortformer with default config (NVIDIA low latency: 1.04s)
-        var config = SortformerConfig.default
+        // Select config (default = NVIDIA low latency ~1.04s). `--config efficient` = chunk_len=25 (~2s, higher throughput).
+        var config: SortformerConfig
+        switch configName {
+        case "efficient":
+            config = .efficientV2_1
+        case "fast", "fastv2_1":
+            config = .fastV2_1
+        case "low", "balanced":
+            config = .balancedV2_1
+        case "high", "highcontext":
+            config = .highContextV2_1
+        default:
+            config = .default
+        }
         var postConfig = DiarizerTimelineConfig.sortformerDefault
         config.debugMode = debugMode
         if let v = predScoreThreshold { config.predScoreThreshold = v }