Support the legacy single-function SpeechDecoder and MultiCodeDecoder assets

.github/workflows/unit-tests.yml

@@ -98,15 +98,16 @@ jobs:
         if: ${{ matrix.run-config['condition'] == true }}
         run: |
           set -o pipefail
-          xcodebuild clean build-for-testing -scheme argmax-oss-swift-Package -destination '${{ matrix.run-config['clean-destination'] }}' | xcbeautify --renderer github-actions
+          # Annotate warnings from a single job only; duplicate check runs each re-post identical annotations on the PR diff
+          xcodebuild clean build-for-testing -scheme argmax-oss-swift-Package -destination '${{ matrix.run-config['clean-destination'] }}' | xcbeautify --renderer ${{ (matrix.run-config['name'] == 'macOS' && inputs.macos-runner == 'macos-26') && 'github-actions' || 'terminal' }}
       - name: Test - ${{ matrix.run-config['name'] }}
         if: ${{ matrix.run-config['condition'] == true }}
         env:
           # xcodebuild strips the TEST_RUNNER_ prefix and forwards HF_TOKEN into the test process
           TEST_RUNNER_HF_TOKEN: ${{ secrets.HF_TOKEN }}
         run: |
           set -o pipefail
-          xcodebuild test -testPlan UnitTestsPlan -scheme argmax-oss-swift-Package -destination '${{ matrix.run-config['test-destination'] }}' | xcbeautify --renderer github-actions
+          xcodebuild test -testPlan UnitTestsPlan -scheme argmax-oss-swift-Package -destination '${{ matrix.run-config['test-destination'] }}' | xcbeautify
       - name: Upload Test Results
         if: failure()
         uses: actions/upload-artifact@ea165f8d65b6e75b540449e92b4886f43607fa02 #v4.6.2

Sources/ArgmaxCore/ModelUtilities.swift

@@ -9,21 +9,40 @@ public struct ModelUtilities {
 
     // MARK: - Model Dimension Introspection
 
-    /// Read a dimension from a model input's multiarray constraint shape.
-    public static func getModelInputDimension(_ model: MLModel?, named: String, position: Int) -> Int? {
+    /// Read the full multiarray constraint shape of a model input.
+    ///
+    /// Returns `nil` when the input is absent or is not a multiarray, which lets
+    /// callers use presence of the shape as a schema probe (e.g. "does this asset
+    /// declare a `qk_mask` input?").
+    public static func getModelInputShape(_ model: MLModel?, named: String) -> [Int]? {
         guard let inputDescription = model?.modelDescription.inputDescriptionsByName[named] else { return nil }
         guard inputDescription.type == .multiArray else { return nil }
         guard let shapeConstraint = inputDescription.multiArrayConstraint else { return nil }
-        let shape = shapeConstraint.shape.map { $0.intValue }
+        return shapeConstraint.shape.map { $0.intValue }
+    }
+
+    /// Read the full multiarray constraint shape of a model output.
+    public static func getModelOutputShape(_ model: MLModel?, named: String) -> [Int]? {
+        guard let outputDescription = model?.modelDescription.outputDescriptionsByName[named] else { return nil }
+        guard outputDescription.type == .multiArray else { return nil }
+        guard let shapeConstraint = outputDescription.multiArrayConstraint else { return nil }
+        return shapeConstraint.shape.map { $0.intValue }
+    }
+
+    /// Read a dimension from a model input's multiarray constraint shape.
+    /// Returns `nil` if the input is absent or `position` is out of range, so that
+    /// probing an unexpected-rank input reports "unknown" instead of trapping.
+    public static func getModelInputDimension(_ model: MLModel?, named: String, position: Int) -> Int? {
+        guard let shape = getModelInputShape(model, named: named) else { return nil }
+        guard position >= 0, position < shape.count else { return nil }
         return shape[position]
     }
 
     /// Read a dimension from a model output's multiarray constraint shape.
+    /// Returns `nil` if the output is absent or `position` is out of range.
     public static func getModelOutputDimension(_ model: MLModel?, named: String, position: Int) -> Int? {
-        guard let inputDescription = model?.modelDescription.outputDescriptionsByName[named] else { return nil }
-        guard inputDescription.type == .multiArray else { return nil }
-        guard let shapeConstraint = inputDescription.multiArrayConstraint else { return nil }
-        let shape = shapeConstraint.shape.map { $0.intValue }
+        guard let shape = getModelOutputShape(model, named: named) else { return nil }
+        guard position >= 0, position < shape.count else { return nil }
         return shape[position]
     }
 

Sources/TTSKit/Protocols.swift

@@ -130,6 +130,11 @@ public protocol SpeechDecoding: MLModelLoading {
     /// SpeechDecoder, 4 for throughput-optimized). Read from the loaded model's
     /// `audio_codes` input shape. Implies `qk_mask` is consumed iff this is > 1.
     var codesPerStep: Int { get }
+    /// `true` when the model's `kv_cache_update_mask` input is rank 3
+    /// (`[1, codesPerStep, maxSeq]`), `false` for the rank-2 `[1, maxSeq]` mask used
+    /// by legacy single-function assets. Determines the mask the host allocates.
+    /// Not derivable from `codesPerStep`: both mask ranks exist at `codesPerStep == 1`.
+    var usesRank3UpdateMask: Bool { get }
 
     // MARK: - Decoding
 
@@ -148,3 +153,9 @@ public protocol SpeechDecoding: MLModelLoading {
         cache: SpeechDecoderCache
     ) async throws -> SpeechDecoderTimedResult
 }
+
+public extension SpeechDecoding {
+    /// Default matching the multifunction asset, so existing conformers that predate
+    /// single-function support keep their behavior without implementing this.
+    var usesRank3UpdateMask: Bool { true }
+}

Sources/TTSKit/Qwen3TTS/Qwen3GenerateTask.swift

@@ -345,7 +345,8 @@ open class Qwen3GenerateTask: @unchecked Sendable, SpeechGenerating {
             maxSeqLength: speechDecoder.kvCacheMaxSequenceLength,
             hiddenDim: speechDecoder.hiddenDim,
             hiddenContextLen: speechDecoder.hiddenContextLen,
-            codesPerStep: codesPerStep
+            codesPerStep: codesPerStep,
+            useRank3UpdateMask: speechDecoder.usesRank3UpdateMask
         )
 
         var generatedTokens: [Int32] = []

Sources/TTSKit/Qwen3TTS/Qwen3Models.swift

@@ -199,8 +199,8 @@ public enum Qwen3SpeechDecoderMode: String, Sendable, CaseIterable {
 /// Selects which MultiCodeDecoder graph expands a talker frame into its 15
 /// residual codes. `.stepped` decodes one position per prediction; `.fused`
 /// decodes the whole frame in one prediction with in-graph sampling.
-/// Loading requires a multifunction asset exposing both functions; legacy
-/// single-function assets fail to load with a function-selection error.
+/// `.fused` requires a multifunction asset; legacy single-function assets are
+/// schema-identical to `stepped` and load fine in that mode.
 @frozen
 public enum Qwen3MultiCodeDecoderMode: String, Sendable, CaseIterable {
     case stepped

Sources/TTSKit/Qwen3TTS/Qwen3MultiCodeDecoder.swift

@@ -29,6 +29,16 @@ struct MLTensorStepResult {
 
 /// Multi-code RVQ decoder backed by a CoreML model.
 ///
+/// Two asset layouts are supported, distinguished at load time by inspecting the
+/// asset rather than by any variant string:
+///
+/// - **Multifunction** (`W8A16-multifunction`): carries a `stepped` graph (one code
+///   position per prediction) and a `fused` graph (whole 15-code frame in one
+///   prediction with in-graph sampling), selected via
+///   `MLModelConfiguration.functionName`.
+/// - **Single-function** (legacy `W8A16`): one graph, schema-identical to the
+///   multifunction `stepped` function, so it drives the same decode path.
+///
 /// Thread safety: mutable state (`model`, dimension properties) is set once during
 /// `loadModel()` and read-only thereafter. `MLModel.prediction()` is thread-safe.
 /// Per-call state is created locally within `generateMultiCodes()` and never stored
@@ -44,42 +54,73 @@ public class Qwen3MultiCodeDecoder: MultiCodeDecoding, @unchecked Sendable {
     public private(set) var codecVocabSize: Int = Qwen3TTSConstants.codecVocabSize
 
     /// Which function of the multifunction asset to load. Set before `loadModel`;
-    /// legacy single-function assets cannot be loaded (function selection fails).
+    /// ignored — beyond a compatibility check — when the asset is single-function.
     public var mode: Qwen3MultiCodeDecoderMode = .stepped
 
+    /// `true` when the loaded asset exposed CoreML functions and one was selected.
+    public private(set) var isMultifunction: Bool = false
+
     public init() {}
 
     public func loadModel(at url: URL, computeUnits: MLComputeUnits, prewarmMode: Bool = false) async throws {
         let modelConfig = MLModelConfiguration()
         modelConfig.computeUnits = computeUnits
-        // The MultiCodeDecoder ships as a multifunction asset carrying both the
-        // `stepped` and `fused` graphs; `mode.functionName` selects one. Function
-        // selection is iOS 18+ / macOS 15+, and the asset requires the same minimum,
-        // so callers on older OS cannot load it.
+        // The tensor decode path is iOS 18+ / macOS 15+, as is CoreML function
+        // selection, so both asset layouts share that floor.
         guard #available(macOS 15.0, iOS 18.0, watchOS 11.0, visionOS 2.0, *) else {
             throw TTSError.modelLoadingFailed(
-                "MultiCodeDecoder requires macOS 15 / iOS 18 (multifunction CoreML model)"
+                "MultiCodeDecoder requires macOS 15 / iOS 18"
             )
         }
-        modelConfig.functionName = mode.functionName
+
+        // Probe the asset for CoreML functions instead of assuming a layout: a
+        // multifunction asset names its graphs, a legacy single-function asset
+        // reports none and must be loaded without a `functionName`.
+        let functionNames = await Self.functionNames(at: url)
+        let selectedFunction: String?
+        if functionNames.isEmpty {
+            // Single-function asset: schema-identical to the multifunction `stepped`
+            // graph, so it can serve `.stepped` but has no in-graph sampling to offer
+            // `.fused`. Say so rather than silently running the stepped path.
+            guard mode == .stepped else {
+                throw TTSError.invalidConfiguration(
+                    "MultiCodeDecoder: \(url.lastPathComponent) is a single-function asset and " +
+                    "only supports \(Qwen3MultiCodeDecoderMode.stepped.rawValue). Use the " +
+                    "multifunction variant ('\(Qwen3VariantDefaults.multiCodeDecoder)') for " +
+                    "\(mode.rawValue)."
+                )
+            }
+            selectedFunction = nil
+        } else {
+            guard functionNames.contains(mode.functionName) else {
+                throw TTSError.modelLoadingFailed(
+                    "MultiCodeDecoder: \(url.lastPathComponent) has no function " +
+                    "'\(mode.functionName)' (available: \(functionNames.joined(separator: ", ")))."
+                )
+            }
+            selectedFunction = mode.functionName
+            modelConfig.functionName = mode.functionName
+        }
+
         let loaded: MLModel
         do {
             loaded = try await MLModel.load(contentsOf: url, configuration: modelConfig)
         } catch {
+            let functionSuffix = selectedFunction.map { " (function '\($0)')" } ?? ""
             throw TTSError.modelLoadingFailed(
-                "MultiCodeDecoder: failed to load function '\(mode.functionName)' from " +
-                "\(url.lastPathComponent). This must be a multifunction CoreML asset " +
-                "with 'stepped' and 'fused' functions, running on macOS 15 / iOS 18 " +
-                "or newer. (\(error.localizedDescription))"
+                "MultiCodeDecoder: failed to load \(url.lastPathComponent)\(functionSuffix) on " +
+                "macOS 15 / iOS 18 or newer. (\(error.localizedDescription))"
             )
         }
 
         // In prewarm mode, compilation is complete - discard to free memory before next model compiles
         guard !prewarmMode else { return }
 
         self.model = loaded
-        // The selected function determines the graph: `fused` samples the whole
-        // frame in-graph, `stepped` decodes one position per prediction.
+        self.isMultifunction = selectedFunction != nil
+        // The selected graph determines the decode path: `fused` samples the whole
+        // frame in-graph, `stepped` (and the legacy asset) decode one position per
+        // prediction.
         self.isFused = mode == .fused
 
         // Detect dimensions from model description
@@ -99,6 +140,19 @@ public class Qwen3MultiCodeDecoder: MultiCodeDecoding, @unchecked Sendable {
         }
     }
 
+    /// CoreML function names exposed by the asset at `url`, or `[]` when it is a
+    /// single-function asset (or cannot be inspected — the subsequent `MLModel.load`
+    /// reports the real failure with better context).
+    @available(macOS 15.0, iOS 18.0, watchOS 11.0, visionOS 2.0, *)
+    private static func functionNames(at url: URL) async -> [String] {
+        do {
+            return try await MLModelAsset(url: url).functionNames
+        } catch {
+            Logging.debug("MultiCodeDecoder: could not read function names from \(url.lastPathComponent): \(error)")
+            return []
+        }
+    }
+
     /// Embedding dimension for `input_embeds`, detected from the model at load time.
     public private(set) var inputEmbedDim: Int = Qwen3TTSConstants.embedDim