Add simple metadata for a couple of models.

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 = "<group>"; };
 		EBB5BA5929426E06003A2A5B /* Downloader.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = Downloader.swift; sourceTree = "<group>"; };
 		EBE3FF4A295DFE2400E921AA /* common.xcconfig */ = {isa = PBXFileReference; lastKnownFileType = text.xcconfig; name = common.xcconfig; path = config/common.xcconfig; sourceTree = "<group>"; };
+		EBE3FF4B295E1EFE00E921AA /* ModelInfo.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = ModelInfo.swift; sourceTree = "<group>"; };
 		EBE4438729488DCA00CDA605 /* README.md */ = {isa = PBXFileReference; lastKnownFileType = net.daringfireball.markdown; path = README.md; sourceTree = "<group>"; };
 		EBE443892948953600CDA605 /* LICENSE */ = {isa = PBXFileReference; lastKnownFileType = text; path = LICENSE; sourceTree = "<group>"; };
 		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 */,

Diffusion/DiffusionApp.swift

@@ -17,4 +17,7 @@ struct DiffusionApp: App {
     }
 }
 
+// A couple of helpers
+
 extension String: Error {}
+let runningOnMac = ProcessInfo.processInfo.isMacCatalystApp

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
+    )
+}

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) {
+    init(model: ModelInfo) {
-        self.url = url
+        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

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 {

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() {