Advanced Settings: ANE (huggingface#28)

* Preparing to allow users to override inference settings.

* ANE setting.

* Do not show advanced settings if ane is not available.

Author: pcuenca

Diffusion-macOS/ControlsView.swift

@@ -56,6 +56,8 @@ struct ControlsView: View {
     @State private var disclosedGuidance = false
     @State private var disclosedSteps = false
     @State private var disclosedSeed = false
+    @State private var disclosedAdvanced = false
+    @State private var useANE = (Settings.shared.userSelectedAttentionVariant ?? ModelInfo.defaultAttention) == .splitEinsum
 
     // TODO: refactor download with similar code in Loading.swift (iOS)
     @State private var stateSubscriber: Cancellable?
@@ -64,28 +66,35 @@ struct ControlsView: View {
 
     // TODO: make this computed, and observable, and easy to read
     @State private var mustShowSafetyCheckerDisclaimer = false
-    
+    @State private var mustShowModelDownloadDisclaimer = false      // When changing advanced settings
+
     @State private var showModelsHelp = false
     @State private var showPromptsHelp = false
     @State private var showGuidanceHelp = false
     @State private var showStepsHelp = false
     @State private var showSeedHelp = false
-    
+    @State private var showAdvancedHelp = false
+
     // Reasonable range for the slider
     let maxSeed: UInt32 = 1000
 
     func updateSafetyCheckerState() {
         mustShowSafetyCheckerDisclaimer = generation.disableSafety && !Settings.shared.safetyCheckerDisclaimerShown
     }
 
+    func updateANEState() {
+        Settings.shared.userSelectedAttentionVariant = useANE ? .splitEinsum : .original
+        modelDidChange(model: Settings.shared.currentModel)
+    }
+    
     func modelDidChange(model: ModelInfo) {
         print("Loading model \(model)")
         Settings.shared.currentModel = model
 
         pipelineLoader?.cancel()
         pipelineState = .downloading(0)
         Task.init {
-            let loader = PipelineLoader(model: model, maxSeed: maxSeed)
+            let loader = PipelineLoader(model: model, variant: Settings.shared.userSelectedAttentionVariant, maxSeed: maxSeed)
             self.pipelineLoader = loader
             stateSubscriber = loader.statePublisher.sink { state in
                 DispatchQueue.main.async {
@@ -114,16 +123,20 @@ struct ControlsView: View {
         }
     }
 
+    func isModelDownloaded(_ model: ModelInfo, variant: AttentionVariant? = nil) -> Bool {
+        PipelineLoader(model: model, variant: variant ?? Settings.shared.userSelectedAttentionVariant).ready
+    }
+    
     func modelLabel(_ model: ModelInfo) -> Text {
-        let downloaded = PipelineLoader(model: model).ready
+        let downloaded = isModelDownloaded(model)
         let prefix = downloaded ? "● " : "◌ "  //"○ "
         return Text(prefix).foregroundColor(downloaded ? .accentColor : .secondary) + Text(model.modelVersion)
     }
 
     var body: some View {
         VStack(alignment: .leading) {
 
-            Label("Adjustments", systemImage: "gearshape.2")
+            Label("Generation Options", systemImage: "gearshape.2")
                 .font(.headline)
                 .fontWeight(.bold)
             Divider()
@@ -217,7 +230,6 @@ struct ControlsView: View {
                             }
                         }.foregroundColor(.secondary)
                     }
-                    Divider()
 
                     DisclosureGroup(isExpanded: $disclosedSteps) {
                         CompactSlider(value: $generation.steps, in: 0...150, step: 5) {
@@ -244,7 +256,6 @@ struct ControlsView: View {
                             }
                         }.foregroundColor(.secondary)
                     }
-                    Divider()
 
                     DisclosureGroup(isExpanded: $disclosedSeed) {
                         let sliderLabel = generation.seed < 0 ? "Random Seed" : "Seed"
@@ -272,6 +283,47 @@ struct ControlsView: View {
                             }
                         }.foregroundColor(.secondary)
                     }
+                    
+                    if hasANE {
+                        Divider()
+                        DisclosureGroup(isExpanded: $disclosedAdvanced) {
+                            HStack {
+                                Toggle("Use Neural Engine", isOn: $useANE).onChange(of: useANE) { value in
+                                    guard let currentModel = ModelInfo.from(modelVersion: model) else { return }
+                                    let variantDownloaded = isModelDownloaded(currentModel, variant: useANE ? .splitEinsum : .original)
+                                    if variantDownloaded {
+                                        updateANEState()
+                                    } else {
+                                        mustShowModelDownloadDisclaimer.toggle()
+                                    }
+                                }
+                                .padding(.leading, 10)
+                                Spacer()
+                            }
+                            .alert("Download Required", isPresented: $mustShowModelDownloadDisclaimer, actions: {
+                                Button("Cancel", role: .destructive) { useANE.toggle() }
+                                Button("Download", role: .cancel) { updateANEState() }
+                            }, message: {
+                                Text("This setting requires a new version of the selected model.")
+                            })
+                        } label: {
+                            HStack {
+                                Label("Advanced", systemImage: "terminal").foregroundColor(.secondary)
+                                Spacer()
+                                if disclosedAdvanced {
+                                    Button {
+                                        showAdvancedHelp.toggle()
+                                    } label: {
+                                        Image(systemName: "info.circle")
+                                    }
+                                    .buttonStyle(.plain)
+                                    .popover(isPresented: $showAdvancedHelp, arrowEdge: .trailing) {
+                                        advancedHelp($showAdvancedHelp)
+                                    }
+                                }
+                            }.foregroundColor(.secondary)
+                        }
+                    }
                 }
             }
             .disclosureGroupStyle(LabelToggleDisclosureGroupStyle())

Diffusion-macOS/Diffusion_macOSApp.swift

@@ -18,3 +18,10 @@ struct Diffusion_macOSApp: App {
 }
 
 let runningOnMac = true
+
+#if canImport(MLCompute)
+import MLCompute
+let hasANE = MLCDevice.ane() != nil
+#else
+let hasANE = false
+#endif

Diffusion-macOS/HelpContent.swift

@@ -123,3 +123,17 @@ func seedHelp(_ showing: Binding<Bool>) -> some View {
          """
     return helpContent(title: "Generation Seed", description: description, showing: showing)
 }
+
+func advancedHelp(_ showing: Binding<Bool>) -> some View {
+    let description =
+         """
+         This section allows you to try different optimization settings.
+         
+         Diffusers will try to select the best configuration for you, but it may not always be optimal \
+         for your computer. You can experiment with these settings to verify the combination that works faster \
+         in your system.
+         
+         Please, note that these settings may trigger downloads of additional model variants.
+         """
+    return helpContent(title: "Advanced Model Settings", description: description, showing: showing)
+}

Diffusion/ModelInfo.swift

@@ -8,6 +8,11 @@
 
 import CoreML
 
+enum AttentionVariant: String {
+    case original
+    case splitEinsum
+}
+
 struct ModelInfo {
     /// Hugging Face model Id that contains .zip archives with compiled Core ML models
     let modelId: String
@@ -19,38 +24,45 @@ struct ModelInfo {
     let originalAttentionSuffix: String
 
     /// Suffix of the archive containing the SPLIT_EINSUM attention variant. Usually something like "split_einsum_compiled"
-    let splitAttentionName: String
+    let splitAttentionSuffix: 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) {
+    init(modelId: String, modelVersion: String, originalAttentionSuffix: String = "original_compiled", splitAttentionSuffix: String = "split_einsum_compiled", supportsEncoder: Bool = false) {
         self.modelId = modelId
         self.modelVersion = modelVersion
         self.originalAttentionSuffix = originalAttentionSuffix
-        self.splitAttentionName = splitAttentionName
+        self.splitAttentionSuffix = splitAttentionSuffix
         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 {
+    static var defaultAttention: AttentionVariant {
+        return runningOnMac ? .original : .splitEinsum
+    }
+    
+    // TODO: heuristics per {model, device}
+    var bestAttention: AttentionVariant {
+        return ModelInfo.defaultAttention
+    }
+    
+    func modelURL(for variant: AttentionVariant) -> 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 suffix: String
+        switch variant {
+        case .original: suffix = originalAttentionSuffix
+        case .splitEinsum: suffix = splitAttentionSuffix
+        }
         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
-    }
-    
+    /// 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 { modelURL(for: bestAttention) }
+        
     var reduceMemory: Bool {
         return !runningOnMac
     }

Diffusion/Pipeline/PipelineLoader.swift

@@ -18,12 +18,14 @@ class PipelineLoader {
     static let models = Path.applicationSupport / "hf-diffusion-models"
 
     let model: ModelInfo
+    let variant: AttentionVariant
     let maxSeed: UInt32
 
     private var downloadSubscriber: Cancellable?
 
-    init(model: ModelInfo, maxSeed: UInt32 = UInt32.max) {
+    init(model: ModelInfo, variant: AttentionVariant? = nil, maxSeed: UInt32 = UInt32.max) {
         self.model = model
+        self.variant = variant ?? model.bestAttention
         self.maxSeed = maxSeed
         state = .undetermined
         setInitialState()
@@ -73,7 +75,7 @@ extension PipelineLoader {
 
 extension PipelineLoader {
     var url: URL {
-        return model.bestURL
+        return model.modelURL(for: variant)
     }
 
     var filename: String {
@@ -95,6 +97,11 @@ extension PipelineLoader {
     var ready: Bool {
         return compiledPath.exists
     }
+    
+    // TODO: measure performance on different devices, disassociate from variant
+    var computeUnits: MLComputeUnits {
+        variant == .original ? .cpuAndGPU : .cpuAndNeuralEngine
+    }
 
     // TODO: maybe receive Progress to add another progress as child
     func prepare() async throws -> Pipeline {
@@ -142,7 +149,7 @@ extension PipelineLoader {
     func load(url: URL) async throws -> StableDiffusionPipeline {
         let beginDate = Date()
         let configuration = MLModelConfiguration()
-        configuration.computeUnits = model.bestComputeUnits
+        configuration.computeUnits = computeUnits
         let pipeline = try StableDiffusionPipeline(resourcesAt: url,
                                                    configuration: configuration,
                                                    disableSafety: false,

Diffusion/State.swift

@@ -73,12 +73,14 @@ class Settings {
     enum Keys: String {
         case model
         case safetyCheckerDisclaimer
+        case variant
     }
 
     private init() {
         defaults.register(defaults: [
             Keys.model.rawValue: ModelInfo.v2Base.modelId,
-            Keys.safetyCheckerDisclaimer.rawValue: false
+            Keys.safetyCheckerDisclaimer.rawValue: false,
+            Keys.variant.rawValue: "- default -"
         ])
     }
 
@@ -100,4 +102,17 @@ class Settings {
             return defaults.bool(forKey: Keys.safetyCheckerDisclaimer.rawValue)
         }
     }
+    
+    /// Returns the option selected by the user, if overridden
+    /// `nil` means: guess best for this {model, device}
+    var userSelectedAttentionVariant: AttentionVariant? {
+        set {
+            // Any String other than the supported ones would cause `get` to return `nil`
+            defaults.set(newValue?.rawValue ?? "- default -", forKey: Keys.variant.rawValue)
+        }
+        get {
+            let current = defaults.string(forKey: Keys.variant.rawValue)
+            return AttentionVariant(rawValue: current ?? "")
+        }
+    }
 }