Resolve compiled assets in diffusion loaders and fail fast on missing files

swift/Sources/CoreAIDiffusionPipeline/Components/CoreAIDiffusionModelFunction.swift

@@ -27,6 +27,10 @@ public actor CoreAIDiffusionModelFunction {
     public func loadResources() async throws {
         guard !isLoaded else { return }
 
+        // Fail fast on a missing asset.
+        guard FileManager.default.fileExists(atPath: modelURL.path) else {
+            throw CoreAIDiffusionError.modelFileNotFound(modelURL)
+        }
         let options = SpecializationOptions(preferredComputeUnitKind: .gpu)
         let loadedModel = try await AIModel(contentsOf: modelURL, options: options)
         guard let fn = try loadedModel.loadFunction(named: "main") else {
@@ -52,6 +56,9 @@ public actor CoreAIDiffusionModelFunction {
     /// The asset is released once this function returns.
     public func hasFunction(named name: String) async throws -> Bool {
         if let model { return model.functionNames.contains(name) }
+        guard FileManager.default.fileExists(atPath: modelURL.path) else {
+            throw CoreAIDiffusionError.modelFileNotFound(modelURL)
+        }
         let options = SpecializationOptions(preferredComputeUnitKind: .gpu)
         let probe = try await AIModel(contentsOf: modelURL, options: options)
         return probe.functionNames.contains(name)
@@ -397,6 +404,7 @@ public actor CoreAIDiffusionModelFunction {
 // MARK: - Errors
 
 public enum CoreAIDiffusionError: Error, LocalizedError {
+    case modelFileNotFound(URL)
     case functionNotFound(String, URL)
     case notLoaded
     case unsupportedInputScalarType(NDArray.ScalarType)
@@ -407,6 +415,8 @@ public enum CoreAIDiffusionError: Error, LocalizedError {
 
     public var errorDescription: String? {
         switch self {
+        case .modelFileNotFound(let url):
+            return "Model asset not found at \(url.path)"
         case .functionNotFound(let name, let url):
             return "Function '\(name)' not found in \(url.lastPathComponent)"
         case .notLoaded:

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

@@ -114,8 +114,12 @@ extension Flux2Pipeline {
             preconditionFailure("auto resolved above")
         }
 
-        let textEncoder = CoreAIDiffusionModelFunction(
-            modelURL: url.appendingPathComponent(textEncoderPath))
+        // Resolve compiled assets; report a missing text encoder up front.
+        let textEncoderURL = ModelBundle.resolveAssetURL(textEncoderPath, in: url)
+        guard FileManager.default.fileExists(atPath: textEncoderURL.path) else {
+            throw PipelineLoadError.missingComponent("text_encoder")
+        }
+        let textEncoder = CoreAIDiffusionModelFunction(modelURL: textEncoderURL)
         let decoder = CoreAIDiffusionModelFunction(
             modelURL: url.appendingPathComponent(decoderName))
 
@@ -131,7 +135,11 @@ extension Flux2Pipeline {
         }
         let encoder: CoreAIDiffusionModelFunction?
         if let name = encoderName {
-            encoder = CoreAIDiffusionModelFunction(modelURL: url.appendingPathComponent(name))
+            // Optional component: nil when absent so supportsImageToImage is accurate.
+            let encoderURL = ModelBundle.resolveAssetURL(name, in: url)
+            encoder = FileManager.default.fileExists(atPath: encoderURL.path)
+                ? CoreAIDiffusionModelFunction(modelURL: encoderURL)
+                : nil
         } else {
             encoder = nil
         }

swift/Sources/CoreAIVideoDiffusionPipeline/Pipelines/WanPipeline.swift

@@ -109,13 +109,13 @@ public struct WanPipeline: VideoPipeline {
 
         self.init(
             transformer: CoreAIDiffusionModelFunction(
-                modelURL: url.appendingPathComponent("Transformer.aimodel")
+                modelURL: ModelBundle.resolveAssetURL("Transformer.aimodel", in: url)
             ),
             textEncoder: CoreAIDiffusionModelFunction(
-                modelURL: url.appendingPathComponent("TextEncoder.aimodel")
+                modelURL: ModelBundle.resolveAssetURL("TextEncoder.aimodel", in: url)
             ),
             decoder: CoreAIDiffusionModelFunction(
-                modelURL: url.appendingPathComponent("VAEDecoder.aimodel")
+                modelURL: ModelBundle.resolveAssetURL("VAEDecoder.aimodel", in: url)
             ),
             tokenizer: tokenizer,
             textDim: textDim,

swift/Tests/CoreAISharedTests/ModelBundleTests.swift

@@ -104,14 +104,34 @@ struct ModelBundleTests {
         #expect(String(describing: error).contains("model.aimodelc"))
     }
 
-    @Test("Pointing at a .aimodel asset throws pointedAtModelAsset")
-    func pointedAtUncompiledAssetThrows() throws {
-        let error = #expect(throws: ModelBundle.BundleError.self) {
-            _ = try ModelBundle(from: "/some/where/model.aimodel")
-        }
-        guard case .pointedAtModelAsset = error else {
-            Issue.record("expected pointedAtModelAsset, got \(String(describing: error))")
-            return
-        }
+    @Test("resolveAssetURL falls back from .aimodel to a compiled .aimodelc")
+    func resolveAssetFallsBackToCompiled() throws {
+        // Mirrors a bundle produced by `coreai-build compile`: metadata.json still
+        // names the .aimodel, but only the compiled .aimodelc exists on disk.
+        let dir = FileManager.default.temporaryDirectory.appending(
+            path: "ModelBundleTests-\(UUID().uuidString)"
+        )
+        try FileManager.default.createDirectory(at: dir, withIntermediateDirectories: true)
+        defer { try? FileManager.default.removeItem(at: dir) }
+        let compiled = dir.appending(path: "TextEncoder.aimodelc")
+        try FileManager.default.createDirectory(at: compiled, withIntermediateDirectories: true)
+
+        let resolved = ModelBundle.resolveAssetURL("TextEncoder.aimodel", in: dir)
+        #expect(resolved == compiled)
+        #expect(FileManager.default.fileExists(atPath: resolved.path))
+    }
+
+    @Test("resolveAssetURL prefers an existing .aimodel over the compiled variant")
+    func resolveAssetPrefersUncompiled() throws {
+        let dir = FileManager.default.temporaryDirectory.appending(
+            path: "ModelBundleTests-\(UUID().uuidString)"
+        )
+        try FileManager.default.createDirectory(at: dir, withIntermediateDirectories: true)
+        defer { try? FileManager.default.removeItem(at: dir) }
+        let asset = dir.appending(path: "TextEncoder.aimodel")
+        try FileManager.default.createDirectory(at: asset, withIntermediateDirectories: true)
+
+        let resolved = ModelBundle.resolveAssetURL("TextEncoder.aimodel", in: dir)
+        #expect(resolved == asset)
     }
 }

swift/Tests/DiffusionPipelineTests/ComponentTests.swift

@@ -51,6 +51,26 @@ struct ComponentTests {
         #expect(count == 100)
     }
 
+    // MARK: - CoreAIDiffusionModelFunction
+
+    @Test("loadResources throws promptly on a missing file rather than hanging")
+    func loadResourcesMissingFileThrows() async {
+        let fn = CoreAIDiffusionModelFunction(
+            modelURL: URL(filePath: "/nonexistent.aimodel"))
+        await #expect(throws: CoreAIDiffusionError.self) {
+            try await fn.loadResources()
+        }
+    }
+
+    @Test("hasFunction throws promptly on a missing file rather than hanging")
+    func hasFunctionMissingFileThrows() async {
+        let fn = CoreAIDiffusionModelFunction(
+            modelURL: URL(filePath: "/nonexistent.aimodel"))
+        await #expect(throws: CoreAIDiffusionError.self) {
+            _ = try await fn.hasFunction(named: "main")
+        }
+    }
+
     // MARK: - CoreAIDenoiser
 
     @Test("Denoiser requires function to be loaded")