diff --git a/swift/Sources/CoreAIDiffusionPipeline/Components/CoreAIDiffusionModelFunction.swift b/swift/Sources/CoreAIDiffusionPipeline/Components/CoreAIDiffusionModelFunction.swift index 9eadd7f6..d48b4964 100644 --- a/swift/Sources/CoreAIDiffusionPipeline/Components/CoreAIDiffusionModelFunction.swift +++ b/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: diff --git a/swift/Sources/CoreAIDiffusionPipeline/Pipelines/Flux2Pipeline+Resources.swift b/swift/Sources/CoreAIDiffusionPipeline/Pipelines/Flux2Pipeline+Resources.swift index 8922f0f0..a282e2d8 100644 --- a/swift/Sources/CoreAIDiffusionPipeline/Pipelines/Flux2Pipeline+Resources.swift +++ b/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 } diff --git a/swift/Sources/CoreAIVideoDiffusionPipeline/Pipelines/WanPipeline.swift b/swift/Sources/CoreAIVideoDiffusionPipeline/Pipelines/WanPipeline.swift index 2b4b6121..65440423 100644 --- a/swift/Sources/CoreAIVideoDiffusionPipeline/Pipelines/WanPipeline.swift +++ b/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, diff --git a/swift/Tests/CoreAISharedTests/ModelBundleTests.swift b/swift/Tests/CoreAISharedTests/ModelBundleTests.swift index 793855de..0281afc9 100644 --- a/swift/Tests/CoreAISharedTests/ModelBundleTests.swift +++ b/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) } } diff --git a/swift/Tests/DiffusionPipelineTests/ComponentTests.swift b/swift/Tests/DiffusionPipelineTests/ComponentTests.swift index ee644240..11764fc8 100644 --- a/swift/Tests/DiffusionPipelineTests/ComponentTests.swift +++ b/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")