diff --git a/Diffusion.xcodeproj/project.pbxproj b/Diffusion.xcodeproj/project.pbxproj index ac5f159..4bbbe0b 100644 --- a/Diffusion.xcodeproj/project.pbxproj +++ b/Diffusion.xcodeproj/project.pbxproj @@ -12,6 +12,7 @@ EBB5BA5829425E17003A2A5B /* Path in Frameworks */ = {isa = PBXBuildFile; productRef = EBB5BA5729425E17003A2A5B /* Path */; }; EBB5BA5A29426E06003A2A5B /* Downloader.swift in Sources */ = {isa = PBXBuildFile; fileRef = EBB5BA5929426E06003A2A5B /* Downloader.swift */; }; EBB5BA5D294504DE003A2A5B /* ZIPFoundation in Frameworks */ = {isa = PBXBuildFile; productRef = EBB5BA5C294504DE003A2A5B /* ZIPFoundation */; }; + EBE3FF4C295E1EFE00E921AA /* ModelInfo.swift in Sources */ = {isa = PBXBuildFile; fileRef = EBE3FF4B295E1EFE00E921AA /* ModelInfo.swift */; }; EBE755C9293E37DD00806B32 /* DiffusionApp.swift in Sources */ = {isa = PBXBuildFile; fileRef = EBE755C8293E37DD00806B32 /* DiffusionApp.swift */; }; EBE755CB293E37DD00806B32 /* TextToImage.swift in Sources */ = {isa = PBXBuildFile; fileRef = EBE755CA293E37DD00806B32 /* TextToImage.swift */; }; EBE755CD293E37DD00806B32 /* Assets.xcassets in Resources */ = {isa = PBXBuildFile; fileRef = EBE755CC293E37DD00806B32 /* Assets.xcassets */; }; @@ -45,6 +46,7 @@ EBB5BA5229425BEE003A2A5B /* PipelineLoader.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = PipelineLoader.swift; sourceTree = ""; }; EBB5BA5929426E06003A2A5B /* Downloader.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = Downloader.swift; sourceTree = ""; }; EBE3FF4A295DFE2400E921AA /* common.xcconfig */ = {isa = PBXFileReference; lastKnownFileType = text.xcconfig; name = common.xcconfig; path = config/common.xcconfig; sourceTree = ""; }; + EBE3FF4B295E1EFE00E921AA /* ModelInfo.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = ModelInfo.swift; sourceTree = ""; }; EBE4438729488DCA00CDA605 /* README.md */ = {isa = PBXFileReference; lastKnownFileType = net.daringfireball.markdown; path = README.md; sourceTree = ""; }; EBE443892948953600CDA605 /* LICENSE */ = {isa = PBXFileReference; lastKnownFileType = text; path = LICENSE; sourceTree = ""; }; EBE755C5293E37DD00806B32 /* Diffusion.app */ = {isa = PBXFileReference; explicitFileType = wrapper.application; includeInIndex = 0; path = Diffusion.app; sourceTree = BUILT_PRODUCTS_DIR; }; @@ -132,6 +134,7 @@ EBE7560A29411A5E00806B32 /* Views */, EBB5BA5929426E06003A2A5B /* Downloader.swift */, EBE755C8293E37DD00806B32 /* DiffusionApp.swift */, + EBE3FF4B295E1EFE00E921AA /* ModelInfo.swift */, EBE755CC293E37DD00806B32 /* Assets.xcassets */, EBE755CE293E37DD00806B32 /* Diffusion.entitlements */, EBE755CF293E37DD00806B32 /* Preview Content */, @@ -330,6 +333,7 @@ EBE75602293E91E200806B32 /* Pipeline.swift in Sources */, EBE755CB293E37DD00806B32 /* TextToImage.swift in Sources */, EBB5BA5A29426E06003A2A5B /* Downloader.swift in Sources */, + EBE3FF4C295E1EFE00E921AA /* ModelInfo.swift in Sources */, EBE756092941178600806B32 /* Loading.swift in Sources */, EBB5BA5329425BEE003A2A5B /* PipelineLoader.swift in Sources */, EBE755C9293E37DD00806B32 /* DiffusionApp.swift in Sources */, diff --git a/Diffusion/DiffusionApp.swift b/Diffusion/DiffusionApp.swift index 265eb1f..a07e3f9 100644 --- a/Diffusion/DiffusionApp.swift +++ b/Diffusion/DiffusionApp.swift @@ -17,4 +17,7 @@ struct DiffusionApp: App { } } +// A couple of helpers + extension String: Error {} +let runningOnMac = ProcessInfo.processInfo.isMacCatalystApp diff --git a/Diffusion/ModelInfo.swift b/Diffusion/ModelInfo.swift new file mode 100644 index 0000000..9365bb2 --- /dev/null +++ b/Diffusion/ModelInfo.swift @@ -0,0 +1,81 @@ +// +// ModelInfo.swift +// Diffusion +// +// Created by Pedro Cuenca on 29/12/22. +// See LICENSE at https://github.com/huggingface/swift-coreml-diffusers/LICENSE +// + +import CoreML + +struct ModelInfo { + /// Hugging Face model Id that contains .zip archives with compiled Core ML models + let modelId: String + + /// Arbitrary string for presentation purposes. Something like "2.1-base" + let modelVersion: String + + /// Suffix of the archive containing the ORIGINAL attention variant. Usually something like "original_compiled" + let originalAttentionSuffix: String + + /// Suffix of the archive containing the SPLIT_EINSUM attention variant. Usually something like "split_einsum_compiled" + let splitAttentionName: String + + /// Whether the archive contains the VAE Encoder (for image to image tasks). Not yet in use. + let supportsEncoder: Bool + + init(modelId: String, modelVersion: String, originalAttentionSuffix: String = "original_compiled", splitAttentionName: String = "split_einsum_compiled", supportsEncoder: Bool = false) { + self.modelId = modelId + self.modelVersion = modelVersion + self.originalAttentionSuffix = originalAttentionSuffix + self.splitAttentionName = splitAttentionName + self.supportsEncoder = supportsEncoder + } +} + +extension ModelInfo { + /// Best variant for the current platform. + /// Currently using `split_einsum` for iOS and `original` for macOS, but could vary depending on model. + var bestURL: URL { + // Pattern: https://huggingface.co/pcuenq/coreml-stable-diffusion/resolve/main/coreml-stable-diffusion-v1-5_original_compiled.zip + let suffix = runningOnMac ? originalAttentionSuffix : splitAttentionName + let repo = modelId.split(separator: "/").last! + return URL(string: "https://huggingface.co/\(modelId)/resolve/main/\(repo)_\(suffix).zip")! + } + + /// Best units for current platform. + /// Currently using `cpuAndNeuralEngine` for iOS and `cpuAndGPU` for macOS, but could vary depending on model. + /// .all works for v1.4, but not for v1.5. + // TODO: measure performance on different devices. + var bestComputeUnits: MLComputeUnits { + return runningOnMac ? .cpuAndGPU : .cpuAndNeuralEngine + } + + var reduceMemory: Bool { + return !runningOnMac + } +} + +extension ModelInfo { + // TODO: repo does not exist yet + static let v14Base = ModelInfo( + modelId: "pcuenq/coreml-stable-diffusion-v1-4", + modelVersion: "1.4" + ) + + static let v15Base = ModelInfo( + modelId: "pcuenq/coreml-stable-diffusion-v1-5", + modelVersion: "1.5" + ) + + static let v2Base = ModelInfo( + modelId: "pcuenq/coreml-stable-diffusion-2-base", + modelVersion: "2-base" + ) + + static let v21Base = ModelInfo( + modelId: "pcuenq/coreml-stable-diffusion-2-1-base", + modelVersion: "2.1-base", + supportsEncoder: true + ) +} diff --git a/Diffusion/Pipeline/PipelineLoader.swift b/Diffusion/Pipeline/PipelineLoader.swift index 08b6766..85fb517 100644 --- a/Diffusion/Pipeline/PipelineLoader.swift +++ b/Diffusion/Pipeline/PipelineLoader.swift @@ -17,11 +17,11 @@ import StableDiffusion class PipelineLoader { static let models = Path.applicationSupport / "hf-diffusion-models" - let url: URL + let model: ModelInfo private var downloadSubscriber: Cancellable? - init(url: URL) { - self.url = url + init(model: ModelInfo) { + self.model = model state = .undetermined setInitialState() } @@ -64,6 +64,10 @@ extension PipelineLoader { } extension PipelineLoader { + var url: URL { + return model.bestURL + } + var filename: String { return url.lastPathComponent } @@ -130,10 +134,11 @@ extension PipelineLoader { func load(url: URL) async throws -> StableDiffusionPipeline { let beginDate = Date() let configuration = MLModelConfiguration() - configuration.computeUnits = .cpuAndGPU // .all works for v1.4, but not for v1.5. TODO: measure performance on different devices + configuration.computeUnits = model.bestComputeUnits let pipeline = try StableDiffusionPipeline(resourcesAt: url, configuration: configuration, - disableSafety: false) + disableSafety: false, + reduceMemory: model.reduceMemory) print("Pipeline loaded in \(Date().timeIntervalSince(beginDate))") state = .loaded return pipeline diff --git a/Diffusion/Views/Loading.swift b/Diffusion/Views/Loading.swift index e7834c7..58e997c 100644 --- a/Diffusion/Views/Loading.swift +++ b/Diffusion/Views/Loading.swift @@ -9,7 +9,7 @@ import SwiftUI import Combine -let PIPELINE_URL = "https://huggingface.co/pcuenq/coreml-stable-diffusion/resolve/main/coreml-stable-diffusion-v1-5_original_compiled.zip" +let model = ModelInfo.v2Base class DiffusionGlobals: ObservableObject { @Published var pipeline: Pipeline? = nil @@ -44,7 +44,7 @@ struct LoadingView: View { .environmentObject(context) .onAppear { Task.init { - let loader = PipelineLoader(url: URL(string: PIPELINE_URL)!) + let loader = PipelineLoader(model: model) stateSubscriber = loader.statePublisher.sink { state in DispatchQueue.main.async { switch state { diff --git a/Diffusion/Views/TextToImage.swift b/Diffusion/Views/TextToImage.swift index aae2c21..779b97b 100644 --- a/Diffusion/Views/TextToImage.swift +++ b/Diffusion/Views/TextToImage.swift @@ -39,7 +39,7 @@ struct ShareButtons: View { var body: some View { let imageView = Image(image, scale: 1, label: Text(name)) - if (ProcessInfo.processInfo.isMacCatalystApp) { + if runningOnMac { HStack { ShareLink(item: imageView, preview: SharePreview(name, image: imageView)) Button() {