diff --git a/Diffusion-macOS/ControlsView.swift b/Diffusion-macOS/ControlsView.swift index ea04b82..9e8ef1c 100644 --- a/Diffusion-macOS/ControlsView.swift +++ b/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,13 +66,15 @@ 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 @@ -78,6 +82,11 @@ struct ControlsView: View { 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 @@ -85,7 +94,7 @@ struct ControlsView: View { 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,8 +123,12 @@ 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) } @@ -123,7 +136,7 @@ struct ControlsView: View { 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()) diff --git a/Diffusion-macOS/Diffusion_macOSApp.swift b/Diffusion-macOS/Diffusion_macOSApp.swift index 623c1c4..1c2c59d 100644 --- a/Diffusion-macOS/Diffusion_macOSApp.swift +++ b/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 diff --git a/Diffusion-macOS/HelpContent.swift b/Diffusion-macOS/HelpContent.swift index 1ef53f5..c60e37a 100644 --- a/Diffusion-macOS/HelpContent.swift +++ b/Diffusion-macOS/HelpContent.swift @@ -123,3 +123,17 @@ func seedHelp(_ showing: Binding) -> some View { """ return helpContent(title: "Generation Seed", description: description, showing: showing) } + +func advancedHelp(_ showing: Binding) -> 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) +} diff --git a/Diffusion/ModelInfo.swift b/Diffusion/ModelInfo.swift index 80ed377..f9865c9 100644 --- a/Diffusion/ModelInfo.swift +++ b/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 } diff --git a/Diffusion/Pipeline/PipelineLoader.swift b/Diffusion/Pipeline/PipelineLoader.swift index 38d628e..3eafff2 100644 --- a/Diffusion/Pipeline/PipelineLoader.swift +++ b/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, diff --git a/Diffusion/State.swift b/Diffusion/State.swift index 89bd706..967b3a9 100644 --- a/Diffusion/State.swift +++ b/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 ?? "") + } + } }