diff --git a/Diffusion-macOS/ControlsView.swift b/Diffusion-macOS/ControlsView.swift index 6469bea..dfdaf82 100644 --- a/Diffusion-macOS/ControlsView.swift +++ b/Diffusion-macOS/ControlsView.swift @@ -57,7 +57,7 @@ struct ControlsView: View { @State private var disclosedSteps = false @State private var disclosedSeed = false @State private var disclosedAdvanced = false - @State private var useANE = 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? @@ -66,7 +66,8 @@ 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 @@ -81,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 @@ -117,8 +123,8 @@ struct ControlsView: View { } } - func isModelDownloaded(_ model: ModelInfo) -> Bool { - PipelineLoader(model: model, variant: Settings.shared.userSelectedAttentionVariant).ready + func isModelDownloaded(_ model: ModelInfo, variant: AttentionVariant? = nil) -> Bool { + PipelineLoader(model: model, variant: variant ?? Settings.shared.userSelectedAttentionVariant).ready } func modelLabel(_ model: ModelInfo) -> Text { @@ -282,10 +288,23 @@ struct ControlsView: View { DisclosureGroup(isExpanded: $disclosedAdvanced) { HStack { Toggle("Use Neural Engine", isOn: $useANE).onChange(of: useANE) { value in - print(value) - }.padding(.leading, 10) + 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) diff --git a/Diffusion/ModelInfo.swift b/Diffusion/ModelInfo.swift index 9b94c9d..f9865c9 100644 --- a/Diffusion/ModelInfo.swift +++ b/Diffusion/ModelInfo.swift @@ -39,10 +39,15 @@ struct ModelInfo { } extension ModelInfo { - var bestAttention: AttentionVariant { + 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: String @@ -57,15 +62,7 @@ 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 { modelURL(for: bestAttention) } - - /// 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 } diff --git a/Diffusion/Pipeline/PipelineLoader.swift b/Diffusion/Pipeline/PipelineLoader.swift index 48597ba..3eafff2 100644 --- a/Diffusion/Pipeline/PipelineLoader.swift +++ b/Diffusion/Pipeline/PipelineLoader.swift @@ -97,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 { @@ -144,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,