Unify asset resolution into ModelBundle.resolveAssetURL (apple#133)

After coreai-build compile produces .aimodelc from .aimodel,
metadata.json still references the original name. ModelBundle now
resolves .aimodel → .aimodelc transparently at the bundle layer.

- Add ModelBundle.resolveAssetURL (single resolution point)
- Wire it into modelURL(for:), verify(), Flux2, SD3, and SD 1.x/2.x
- Remove PreparedModel.resolveCoreAIModelURL (dead after bundle-layer resolution)
- Remove redundant resolution call from EngineFactory.createEngine

Author: stikves

swift/Sources/CoreAIDiffusionPipeline/Pipelines/Flux2Pipeline+Resources.swift

@@ -4,6 +4,7 @@
 // be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause
 
 import CoreAI
+import CoreAIShared
 import Foundation
 import Tokenizers
 
@@ -114,15 +115,9 @@ extension Flux2Pipeline {
 
     /// Resolve an asset name to a filename, checking for .aimodel or .aimodelc.
     private static func resolveAsset(at url: URL, name: String) -> String? {
-        let fm = FileManager.default
-        let aimodel = "\(name).aimodel"
-        let aimodelc = "\(name).aimodelc"
-        if fm.fileExists(atPath: url.appendingPathComponent(aimodel).path) {
-            return aimodel
-        } else if fm.fileExists(atPath: url.appendingPathComponent(aimodelc).path) {
-            return aimodelc
-        }
-        return nil
+        let resolved = ModelBundle.resolveAssetURL("\(name).aimodel", in: url)
+        guard FileManager.default.fileExists(atPath: resolved.path) else { return nil }
+        return resolved.lastPathComponent
     }
 
     /// Probe available assets and pick the highest quality mode.

swift/Sources/CoreAIDiffusionPipeline/Pipelines/PipelineDescriptor+CoreAI.swift

@@ -4,6 +4,7 @@
 // be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause
 
 import CoreAI
+import CoreAIShared
 import Foundation
 
 /// Loaded diffusion pipeline components backed by Core AI model functions.
@@ -51,14 +52,14 @@ extension PipelineDescriptor {
 
         // Create model functions
         let unetFunction = CoreAIDiffusionModelFunction(
-            modelURL: baseURL.appendingPathComponent(unetPath))
+            modelURL: ModelBundle.resolveAssetURL(unetPath, in: baseURL))
         let decoderFunction = CoreAIDiffusionModelFunction(
-            modelURL: baseURL.appendingPathComponent(decoderPath))
+            modelURL: ModelBundle.resolveAssetURL(decoderPath, in: baseURL))
 
         let encoderFunction: CoreAIDiffusionModelFunction?
         if let encoderPath = components.vaeEncoder {
             encoderFunction = CoreAIDiffusionModelFunction(
-                modelURL: baseURL.appendingPathComponent(encoderPath))
+                modelURL: ModelBundle.resolveAssetURL(encoderPath, in: baseURL))
         } else {
             encoderFunction = nil
         }
@@ -108,7 +109,7 @@ extension PipelineDescriptor {
         let textEncoderFunction: CoreAIDiffusionModelFunction
         if let tePath = components.textEncoder {
             textEncoderFunction = CoreAIDiffusionModelFunction(
-                modelURL: baseURL.appendingPathComponent(tePath))
+                modelURL: ModelBundle.resolveAssetURL(tePath, in: baseURL))
         } else {
             throw PipelineLoadError.missingComponent("text_encoder")
         }

swift/Sources/CoreAIDiffusionPipeline/Pipelines/SD3Pipeline+Resources.swift

@@ -4,6 +4,7 @@
 // be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause
 
 import CoreAI
+import CoreAIShared
 import Foundation
 
 extension SD3Pipeline {
@@ -29,13 +30,13 @@ extension SD3Pipeline {
         }
 
         let transformer = CoreAIDiffusionModelFunction(
-            modelURL: url.appendingPathComponent(transformerPath))
+            modelURL: ModelBundle.resolveAssetURL(transformerPath, in: url))
         let textEncoder = CoreAIDiffusionModelFunction(
-            modelURL: url.appendingPathComponent(textEncoderPath))
+            modelURL: ModelBundle.resolveAssetURL(textEncoderPath, in: url))
         let textEncoder2 = CoreAIDiffusionModelFunction(
-            modelURL: url.appendingPathComponent(textEncoder2Path))
+            modelURL: ModelBundle.resolveAssetURL(textEncoder2Path, in: url))
         let decoder = CoreAIDiffusionModelFunction(
-            modelURL: url.appendingPathComponent(decoderPath))
+            modelURL: ModelBundle.resolveAssetURL(decoderPath, in: url))
 
         let tokenizer = try Self.loadBPETokenizer(
             at: url.appendingPathComponent("tokenizer"))

swift/Sources/CoreAILanguageModels/InferenceEngines/EngineFactory.swift

@@ -25,7 +25,7 @@ public struct EngineFactory: Sendable {
     ///
     /// - Parameters:
     ///   - config: The JSON model configuration as raw data.
-    ///   - modelURL: The location of the model asset on disk.
+    ///   - modelURL: The location of the CoreAI model asset (`.aimodel` or `.aimodelc`).
     ///   - options: The engine options, including an optional variant override and KV cache
     ///     settings. Defaults to a value with auto-detection enabled and the `.auto` KV cache
     ///     strategy.
@@ -45,23 +45,20 @@ public struct EngineFactory: Sendable {
             CLILogger.log("  - kvCacheSize: \(size) (override)")
         }
 
-        // Step 2: Resolve model URL for Core AI path
-        let coreAIModelURL = PreparedModel.resolveCoreAIModelURL(from: modelURL)
-
-        // Step 3: Prepare model asset via Core AI
-        let preparedModel = try await PreparedModel.prepare(at: coreAIModelURL)
+        // Step 2: Prepare model asset via Core AI
+        let preparedModel = try await PreparedModel.prepare(at: modelURL)
 
         CLILogger.log("  - structure: \(preparedModel.structure.description)")
 
-        // Step 4: Resolve variant with structure detection
+        // Step 3: Resolve variant with structure detection
         let variant = try resolveVariant(
             override: options.variant,
             detectedStructure: preparedModel.structure
         )
 
         CLILogger.log("  - resolved variant: \(variant.rawValue)")
 
-        // Step 5: Instantiate the appropriate engine
+        // Step 4: Instantiate the appropriate engine
         return try await selectEngine(
             variant: variant,
             config: parsedConfig,

swift/Sources/CoreAIShared/Bundle/ModelBundle.swift

@@ -48,9 +48,10 @@ public struct ModelBundle: Sendable {
     }
 
     /// Resolve a component's URL within the bundle by role key.
+    /// Falls back to `.aimodelc` if the declared `.aimodel` path doesn't exist on disk.
     public func modelURL(for key: String) -> URL? {
         guard let path = assets[key] else { return nil }
-        return bundlePath.appending(path: path)
+        return Self.resolveAssetURL(path, in: bundlePath)
     }
 
     /// Required-component variant — throws `BundleError.missingField` if absent.
@@ -61,13 +62,27 @@ public struct ModelBundle: Sendable {
         return url
     }
 
+    /// Resolve an asset path against a directory, falling back from `.aimodel` to `.aimodelc`.
+    ///
+    /// When `coreai-build compile` produces a compiled `.aimodelc` from a source `.aimodel`,
+    /// metadata.json still references the original name. This finds the compiled variant
+    /// so users don't need to hand-edit metadata.json after compilation.
+    public static func resolveAssetURL(_ path: String, in directory: URL) -> URL {
+        let url = directory.appending(path: path)
+        if FileManager.default.fileExists(atPath: url.path) { return url }
+        if path.hasSuffix(".aimodel") {
+            let compiled = directory.appending(path: path + "c")
+            if FileManager.default.fileExists(atPath: compiled.path) { return compiled }
+        }
+        return url
+    }
+
     /// Verify all declared assets exist on disk. Throws `BundleError.missingAsset`
     /// with guidance if a component is missing (e.g. after manual compilation).
     public func verify() throws {
-        let fm = FileManager.default
         for (key, filename) in assets {
-            let url = bundlePath.appending(path: filename)
-            if !fm.fileExists(atPath: url.path) {
+            let url = Self.resolveAssetURL(filename, in: bundlePath)
+            if !FileManager.default.fileExists(atPath: url.path) {
                 throw BundleError.missingAsset(key: key, path: url)
             }
         }

swift/Sources/CoreAIShared/Runtime/ModelStructure.swift

@@ -104,33 +104,6 @@ public struct PreparedModel: Sendable {
     /// Detected model structure (chunked/static vs dynamic).
     public let structure: ModelStructure
 
-    // MARK: - Core AI Model URL Resolution
-
-    /// If `url` is already `.aimodel`, returns it unchanged. Otherwise looks for
-    /// a sibling `.aimodel` directory with the same base name.
-    public static func resolveCoreAIModelURL(from url: URL) -> URL {
-        let ext = url.pathExtension
-
-        // Already a Core AI format
-        if ext == "aimodel" {
-            return url
-        }
-
-        // Check for sibling Core AI model directory
-        let parentDir = url.deletingLastPathComponent()
-        let baseName = url.deletingPathExtension().lastPathComponent
-
-        let candidate = parentDir.appendingPathComponent("\(baseName).aimodel")
-        if FileManager.default.fileExists(atPath: candidate.path) {
-            CLILogger.log(
-                "  - Resolved CoreAI model path: \(url.lastPathComponent) → \(candidate.lastPathComponent)")
-            return candidate
-        }
-
-        // Fall through to original URL (AIModel may still handle it)
-        return url
-    }
-
     // MARK: - Cache Inspection
 
     /// File extensions that identify a Core AI model asset (source or compiled).
@@ -175,8 +148,7 @@ public struct PreparedModel: Sendable {
     public static func clearCache(at url: URL) throws -> [URL] {
         let assetURLs = try modelAssetURLs(at: url)
         for assetURL in assetURLs {
-            let coreaiURL = resolveCoreAIModelURL(from: assetURL)
-            try AIModelCache.default.deleteEntries(for: coreaiURL)
+            try AIModelCache.default.deleteEntries(for: assetURL)
         }
         return assetURLs
     }
@@ -192,9 +164,8 @@ public struct PreparedModel: Sendable {
     ///   the ``isCached(at:)`` overload; callers that load via `AIModel(contentsOf:)` or a custom
     ///   `SpecializationOptions` must pass the same value here.
     public static func isCached(at url: URL, options: SpecializationOptions) -> Bool {
-        let coreaiURL = resolveCoreAIModelURL(from: url)
         do {
-            return try AIModelCache.default.model(for: coreaiURL, options: options) != nil
+            return try AIModelCache.default.model(for: url, options: options) != nil
         } catch {
             return false
         }
@@ -206,9 +177,8 @@ public struct PreparedModel: Sendable {
     /// Use this only for models loaded through ``prepare(at:)``. For other loaders, use
     /// ``isCached(at:options:)`` with the matching options.
     public static func isCached(at url: URL) -> Bool {
-        let coreaiURL = resolveCoreAIModelURL(from: url)
-        let options = probeStructure(at: coreaiURL).specializationOptions
-        return isCached(at: coreaiURL, options: options)
+        let options = probeStructure(at: url).specializationOptions
+        return isCached(at: url, options: options)
     }
 
     // MARK: - Asset Preparation
@@ -219,7 +189,7 @@ public struct PreparedModel: Sendable {
     /// dynamic models prefer GPU with frequent reshapes; chunked-static models prefer Neural Engine.
     ///
     /// - Parameters:
-    ///   - url: URL to the model asset (`.aimodel` bundle)
+    ///   - url: URL to the model asset (`.aimodel` or `.aimodelc` bundle)
     /// - Returns: Prepared asset with compiled library and detected structure
     /// - Throws: Error from `AIModel` if loading or specialization fails
     public static func prepare(

swift/Sources/Tools/llm-runner/LLMRunnerMain.swift

@@ -376,7 +376,7 @@ struct LLMRunner: AsyncParsableCommand, Sendable {
             cacheHit = PreparedModel.isCached(at: languageModelURL)
         }
 
-        let assetLabel = try modelAssetTypeLabel(for: bundle.modelAssetPath)
+        let assetLabel = try modelAssetTypeLabel(for: languageModelURL.pathExtension)
         if !CLILogger.isVerbose {
             print("\n⏳ Preparing AI asset from \(assetLabel)...", terminator: "")
             fflush(stdout)
@@ -1075,8 +1075,8 @@ struct LLMRunner: AsyncParsableCommand, Sendable {
 
     // MARK: - Asset Type Label
 
-    private func modelAssetTypeLabel(for path: String) throws -> String {
-        switch URL(fileURLWithPath: path).pathExtension.lowercased() {
+    private func modelAssetTypeLabel(for ext: String) throws -> String {
+        switch ext.lowercased() {
         case "aimodelc": return "compiled"
         case "aimodel": return "source"
         default: