diff --git a/Diffusion-macOS/Capabilities.swift b/Diffusion-macOS/Capabilities.swift index 5a419e2..c5ae1bd 100644 --- a/Diffusion-macOS/Capabilities.swift +++ b/Diffusion-macOS/Capabilities.swift @@ -9,6 +9,18 @@ import Foundation let runningOnMac = true +let deviceHas6GBOrMore = true +let deviceHas8GBOrMore = true +let BENCHMARK = false + +let deviceSupportsQuantization = { + if #available(macOS 14, *) { + return true + } else { + return false + } +}() + #if canImport(MLCompute) import MLCompute diff --git a/Diffusion-macOS/ControlsView.swift b/Diffusion-macOS/ControlsView.swift index a36da24..63a6ef6 100644 --- a/Diffusion-macOS/ControlsView.swift +++ b/Diffusion-macOS/ControlsView.swift @@ -10,6 +10,7 @@ import Combine import SwiftUI import CompactSlider +/// Track a StableDiffusion Pipeline's readiness. This include actively downloading from the internet, uncompressing the downloaded zip file, actively loading into memory, ready to use or an Error state. enum PipelineState { case downloading(Double) case uncompressing @@ -55,6 +56,7 @@ struct ControlsView: View { @State private var disclosedPrompt = true @State private var disclosedGuidance = false @State private var disclosedSteps = false + @State private var disclosedPreview = false @State private var disclosedSeed = false @State private var disclosedAdvanced = false @@ -71,32 +73,50 @@ struct ControlsView: View { @State private var showPromptsHelp = false @State private var showGuidanceHelp = false @State private var showStepsHelp = false + @State private var showPreviewHelp = false @State private var showSeedHelp = false @State private var showAdvancedHelp = false + @State private var positiveTokenCount: Int = 0 + @State private var negativeTokenCount: Int = 0 - // Reasonable range for the slider - let maxSeed: UInt32 = 1000 - - func updateSafetyCheckerState() { + let maxSeed: UInt32 = UInt32.max + private var textFieldLabelSeed: String { generation.seed < 1 ? "Random Seed" : "Seed" } + + var modelFilename: String? { + guard let pipelineLoader = pipelineLoader else { return nil } + let selectedURL = pipelineLoader.compiledURL + guard FileManager.default.fileExists(atPath: selectedURL.path) else { return nil } + return selectedURL.path + } + + fileprivate func updateSafetyCheckerState() { mustShowSafetyCheckerDisclaimer = generation.disableSafety && !Settings.shared.safetyCheckerDisclaimerShown } - func updateComputeUnitsState() { + fileprivate func updateComputeUnitsState() { Settings.shared.userSelectedComputeUnits = generation.computeUnits modelDidChange(model: Settings.shared.currentModel) } - func resetComputeUnitsState() { + fileprivate func resetComputeUnitsState() { generation.computeUnits = Settings.shared.userSelectedComputeUnits ?? ModelInfo.defaultComputeUnits } - - func modelDidChange(model: ModelInfo) { + + fileprivate func modelDidChange(model: ModelInfo) { guard pipelineLoader?.model != model || pipelineLoader?.computeUnits != generation.computeUnits else { print("Reusing same model \(model) with units \(generation.computeUnits)") return } - print("Loading model \(model)") + if !model.supportsNeuralEngine && generation.computeUnits == .cpuAndNeuralEngine { + // Reset compute units to GPU if Neural Engine is not supported + Settings.shared.userSelectedComputeUnits = .cpuAndGPU + resetComputeUnitsState() + print("Neural Engine not supported for model \(model), switching to GPU") + } else { + resetComputeUnitsState() + } + Settings.shared.currentModel = model pipelineLoader?.cancel() @@ -108,7 +128,6 @@ struct ControlsView: View { DispatchQueue.main.async { switch state { case .downloading(let progress): - print("\(loader.model.modelVersion): \(progress)") pipelineState = .downloading(progress) case .uncompressing: pipelineState = .uncompressing @@ -131,21 +150,33 @@ struct ControlsView: View { } } - func isModelDownloaded(_ model: ModelInfo, computeUnits: ComputeUnits? = nil) -> Bool { + fileprivate func isModelDownloaded(_ model: ModelInfo, computeUnits: ComputeUnits? = nil) -> Bool { PipelineLoader(model: model, computeUnits: computeUnits ?? generation.computeUnits).ready } - func modelLabel(_ model: ModelInfo) -> Text { + fileprivate func modelLabel(_ model: ModelInfo) -> Text { let downloaded = isModelDownloaded(model) let prefix = downloaded ? "● " : "◌ " //"○ " return Text(prefix).foregroundColor(downloaded ? .accentColor : .secondary) + Text(model.modelVersion) } - var modelFilename: String? { - guard let pipelineLoader = pipelineLoader else { return nil } - let selectedPath = pipelineLoader.compiledPath - guard selectedPath.exists else { return nil } - return selectedPath.string + fileprivate func prompts() -> some View { + VStack { + Spacer() + PromptTextField(text: $generation.positivePrompt, isPositivePrompt: true, model: $model) + .onChange(of: generation.positivePrompt) { prompt in + Settings.shared.prompt = prompt + } + .padding(.top, 5) + Spacer() + PromptTextField(text: $generation.negativePrompt, isPositivePrompt: false, model: $model) + .onChange(of: generation.negativePrompt) { negativePrompt in + Settings.shared.negativePrompt = negativePrompt + } + .padding(.bottom, 5) + Spacer() + } + .frame(maxHeight: .infinity) } var body: some View { @@ -168,7 +199,8 @@ struct ControlsView: View { } .onChange(of: model) { selection in guard selection != revealOption else { - NSWorkspace.shared.selectFile(modelFilename, inFileViewerRootedAtPath: PipelineLoader.models.string) + // The reveal option has been requested - open the models folder in Finder + NSWorkspace.shared.selectFile(modelFilename, inFileViewerRootedAtPath: PipelineLoader.models.path) model = Settings.shared.currentModel.modelVersion return } @@ -197,13 +229,7 @@ struct ControlsView: View { DisclosureGroup(isExpanded: $disclosedPrompt) { Group { - TextField("Positive prompt", text: $generation.positivePrompt, - axis: .vertical).lineLimit(5) - .textFieldStyle(.squareBorder) - .listRowInsets(EdgeInsets(top: 0, leading: -20, bottom: 0, trailing: 20)) - TextField("Negative prompt", text: $generation.negativePrompt, - axis: .vertical).lineLimit(5) - .textFieldStyle(.squareBorder) + prompts() }.padding(.leading, 10) } label: { HStack { @@ -231,7 +257,11 @@ struct ControlsView: View { Text("Guidance Scale") Spacer() Text(guidanceScaleValue) - }.padding(.leading, 10) + } + .onChange(of: generation.guidanceScale) { guidanceScale in + Settings.shared.guidanceScale = guidanceScale + } + .padding(.leading, 10) } label: { HStack { Label("Guidance Scale", systemImage: "scalemass").foregroundColor(.secondary) @@ -254,11 +284,15 @@ struct ControlsView: View { } DisclosureGroup(isExpanded: $disclosedSteps) { - CompactSlider(value: $generation.steps, in: 0...150, step: 5) { + CompactSlider(value: $generation.steps, in: 1...150, step: 1) { Text("Steps") Spacer() Text("\(Int(generation.steps))") - }.padding(.leading, 10) + } + .onChange(of: generation.steps) { steps in + Settings.shared.stepCount = steps + } + .padding(.leading, 10) } label: { HStack { Label("Step count", systemImage: "square.3.layers.3d.down.left").foregroundColor(.secondary) @@ -278,17 +312,43 @@ struct ControlsView: View { } }.foregroundColor(.secondary) } - - DisclosureGroup(isExpanded: $disclosedSeed) { - let sliderLabel = generation.seed < 0 ? "Random Seed" : "Seed" - CompactSlider(value: $generation.seed, in: -1...Double(maxSeed), step: 1) { - Text(sliderLabel) + + DisclosureGroup(isExpanded: $disclosedPreview) { + CompactSlider(value: $generation.previews, in: 0...25, step: 1) { + Text("Previews") Spacer() - Text("\(Int(generation.seed))") - }.padding(.leading, 10) + Text("\(Int(generation.previews))") + } + .onChange(of: generation.previews) { previews in + Settings.shared.previewCount = previews + } + .padding(.leading, 10) } label: { HStack { - Label("Seed", systemImage: "leaf").foregroundColor(.secondary) + Label("Preview count", systemImage: "eye.square").foregroundColor(.secondary) + Spacer() + if disclosedPreview { + Button { + showPreviewHelp.toggle() + } label: { + Image(systemName: "info.circle") + } + .buttonStyle(.plain) + .popover(isPresented: $showPreviewHelp, arrowEdge: .trailing) { + previewHelp($showPreviewHelp) + } + } else { + Text("\(Int(generation.previews))") + } + }.foregroundColor(.secondary) + } + + DisclosureGroup(isExpanded: $disclosedSeed) { + discloseSeedContent() + .padding(.leading, 10) + } label: { + HStack { + Label(textFieldLabelSeed, systemImage: "leaf").foregroundColor(.secondary) Spacer() if disclosedSeed { Button { @@ -301,24 +361,32 @@ struct ControlsView: View { seedHelp($showSeedHelp) } } else { - Text("\(Int(generation.seed))") + Text(generation.seed.formatted(.number.grouping(.never))) } - }.foregroundColor(.secondary) + } + .foregroundColor(.secondary) } - + if Capabilities.hasANE { Divider() + let isNeuralEngineDisabled = !(ModelInfo.from(modelVersion: model)?.supportsNeuralEngine ?? true) DisclosureGroup(isExpanded: $disclosedAdvanced) { HStack { Picker(selection: $generation.computeUnits, label: Text("Use")) { Text("GPU").tag(ComputeUnits.cpuAndGPU) - Text("Neural Engine").tag(ComputeUnits.cpuAndNeuralEngine) + Text("Neural Engine\(isNeuralEngineDisabled ? " (unavailable)" : "")") + .foregroundColor(isNeuralEngineDisabled ? .secondary : .primary) + .tag(ComputeUnits.cpuAndNeuralEngine) Text("GPU and Neural Engine").tag(ComputeUnits.all) }.pickerStyle(.radioGroup).padding(.leading) Spacer() } .onChange(of: generation.computeUnits) { units in guard let currentModel = ModelInfo.from(modelVersion: model) else { return } + if isNeuralEngineDisabled && units == .cpuAndNeuralEngine { + resetComputeUnitsState() + return + } let variantDownloaded = isModelDownloaded(currentModel, computeUnits: units) if variantDownloaded { updateComputeUnitsState() @@ -384,8 +452,45 @@ struct ControlsView: View { } .padding() .onAppear { - print(PipelineLoader.models) modelDidChange(model: ModelInfo.from(modelVersion: model) ?? ModelInfo.v2Base) } } + + fileprivate func discloseSeedContent() -> some View { + let seedBinding = Binding( + get: { + String(generation.seed) + }, + set: { newValue in + if let seed = UInt32(newValue) { + generation.seed = seed + Settings.shared.seed = seed + } else { + generation.seed = 0 + Settings.shared.seed = 0 + } + } + ) + + return HStack { + TextField("", text: seedBinding) + .multilineTextAlignment(.trailing) + .onChange(of: seedBinding.wrappedValue, perform: { newValue in + if let seed = UInt32(newValue) { + generation.seed = seed + Settings.shared.seed = seed + } else { + generation.seed = 0 + Settings.shared.seed = 0 + } + }) + .onReceive(Just(seedBinding.wrappedValue)) { newValue in + let filtered = newValue.filter { "0123456789".contains($0) } + if filtered != newValue { + seedBinding.wrappedValue = filtered + } + } + Stepper("", value: $generation.seed, in: 0...UInt32.max) + } + } } diff --git a/Diffusion-macOS/DiffusionImage+macOS.swift b/Diffusion-macOS/DiffusionImage+macOS.swift new file mode 100644 index 0000000..b3be834 --- /dev/null +++ b/Diffusion-macOS/DiffusionImage+macOS.swift @@ -0,0 +1,112 @@ +// +// DiffusionImage+macOS.swift +// Diffusion-macOS +// +// Created by Dolmere and Pedro Cuenca on 30/07/2023. +// + +import SwiftUI +import UniformTypeIdentifiers + +extension DiffusionImage { + + /// Instance func to place the generated image on the file system and return the `fileURL` where it is stored. + func save(cgImage: CGImage, filename: String?) -> URL? { + + let nsImage = NSImage(cgImage: cgImage, size: NSSize(width: cgImage.width, height: cgImage.height)) + + + let appSupportURL = Settings.shared.tempStorageURL() + let fn = filename ?? "diffusion_generated_image" + let fileURL = appSupportURL + .appendingPathComponent(fn) + .appendingPathExtension("png") + + // Save the image as a temporary file + if let tiffData = nsImage.tiffRepresentation, + let bitmap = NSBitmapImageRep(data: tiffData), + let pngData = bitmap.representation(using: .png, properties: [:]) { + do { + try pngData.write(to: fileURL) + return fileURL + } catch { + print("Error saving image to temporary file: \(error)") + } + } + return nil + } + + /// Returns a `Data` representation of this generated image in PNG format or nil if there is an error converting the image data. + func pngRepresentation() -> Data? { + let bitmapRep = NSBitmapImageRep(cgImage: cgImage) + return bitmapRep.representation(using: .png, properties: [:]) + } +} + +extension DiffusionImage: NSItemProviderWriting { + + // MARK: - NSItemProviderWriting + + static var writableTypeIdentifiersForItemProvider: [String] { + return [UTType.data.identifier, UTType.png.identifier, UTType.fileURL.identifier] + } + + func itemProviderVisibilityForRepresentation(withTypeIdentifier typeIdentifier: String) -> NSItemProviderRepresentationVisibility { + return .all + } + + func itemProviderRepresentation(forTypeIdentifier typeIdentifier: String) throws -> NSItemProvider { + print("itemProviderRepresentation(forTypeIdentifier") + print(typeIdentifier) + let data = try NSKeyedArchiver.archivedData(withRootObject: self, requiringSecureCoding: true) + let itemProvider = NSItemProvider() + itemProvider.registerDataRepresentation(forTypeIdentifier: typeIdentifier, visibility: NSItemProviderRepresentationVisibility.all) { completion in + completion(data, nil) + return nil + } + return itemProvider + } + + func loadData(withTypeIdentifier typeIdentifier: String, forItemProviderCompletionHandler completionHandler: @escaping @Sendable (Data?, Error?) -> Void) -> Progress? { + if typeIdentifier == NSPasteboard.PasteboardType.fileURL.rawValue { + let data = fileURL.dataRepresentation + completionHandler(data, nil) + } else if typeIdentifier == UTType.png.identifier { + let data = pngRepresentation() + completionHandler(data, nil) + } else { + // Indicate that the specified typeIdentifier is not supported + let error = NSError(domain: "com.huggingface.diffusion", code: 0, userInfo: [NSLocalizedDescriptionKey: "Unsupported typeIdentifier"]) + completionHandler(nil, error) + } + return nil + } + +} + +extension DiffusionImage: NSPasteboardWriting { + + // MARK: - NSPasteboardWriting + + func writableTypes(for pasteboard: NSPasteboard) -> [NSPasteboard.PasteboardType] { + return [ + NSPasteboard.PasteboardType.fileURL, + NSPasteboard.PasteboardType(rawValue: UTType.png.identifier) + ] + } + + func pasteboardPropertyList(forType type: NSPasteboard.PasteboardType) -> Any? { + if type == NSPasteboard.PasteboardType.fileURL { + + // Return the file's data' representation + return fileURL.dataRepresentation + + } else if type.rawValue == UTType.png.identifier { + + // Return a PNG data representation + return pngRepresentation() + } + + return nil + } +} diff --git a/Diffusion-macOS/GeneratedImageView.swift b/Diffusion-macOS/GeneratedImageView.swift index 47d2271..de1e30f 100644 --- a/Diffusion-macOS/GeneratedImageView.swift +++ b/Diffusion-macOS/GeneratedImageView.swift @@ -10,7 +10,7 @@ import SwiftUI struct GeneratedImageView: View { @EnvironmentObject var generation: GenerationContext - + var body: some View { switch generation.state { case .startup: return AnyView(Image("placeholder").resizable()) @@ -19,26 +19,47 @@ struct GeneratedImageView: View { // The first time it takes a little bit before generation starts return AnyView(ProgressView()) } + let step = Int(progress.step) + 1 let fraction = Double(step) / Double(progress.stepCount) let label = "Step \(step) of \(progress.stepCount)" - return AnyView(HStack { - ProgressView(label, value: fraction, total: 1).padding() - Button { - generation.cancelGeneration() - } label: { - Image(systemName: "x.circle.fill").foregroundColor(.gray) + + return AnyView(VStack { + Group { + if let safeImage = generation.previewImage { + Image(safeImage, scale: 1, label: Text("generated")) + .resizable() + .clipShape(RoundedRectangle(cornerRadius: 20)) + } + } + HStack { + ProgressView(label, value: fraction, total: 1).padding() + Button { + generation.cancelGeneration() + } label: { + Image(systemName: "x.circle.fill").foregroundColor(.gray) + } + .buttonStyle(.plain) } - .buttonStyle(.plain) }) case .complete(_, let image, _, _): guard let theImage = image else { return AnyView(Image(systemName: "exclamationmark.triangle").resizable()) } - - return AnyView(Image(theImage, scale: 1, label: Text("generated")) - .resizable() - .clipShape(RoundedRectangle(cornerRadius: 20)) + + return AnyView( + Image(theImage, scale: 1, label: Text("generated")) + .resizable() + .clipShape(RoundedRectangle(cornerRadius: 20)) + .contextMenu { + Button { + NSPasteboard.general.clearContents() + let nsimage = NSImage(cgImage: theImage, size: NSSize(width: theImage.width, height: theImage.height)) + NSPasteboard.general.writeObjects([nsimage]) + } label: { + Text("Copy Photo") + } + } ) case .failed(_): return AnyView(Image(systemName: "exclamationmark.triangle").resizable()) diff --git a/Diffusion-macOS/HelpContent.swift b/Diffusion-macOS/HelpContent.swift index c60e37a..e21b662 100644 --- a/Diffusion-macOS/HelpContent.swift +++ b/Diffusion-macOS/HelpContent.swift @@ -110,6 +110,25 @@ func stepsHelp(_ showing: Binding) -> some View { return helpContent(title: "Inference Steps", description: description, showing: showing) } +func previewHelp(_ showing: Binding) -> some View { + let description = + """ + This number controls how many previews to display throughout the image generation process. + + Using more previews can be useful if you want more visibility into how \ + generation is progressing. + + However, computing each preview takes some time and can slow down \ + generation. If the process is too slow you can reduce the preview count, \ + which will result in less visibility of intermediate steps during generation. + + You can try different values to see what works best for your hardware. + + For the absolute fastest generation times, use 0 previews. + """ + return helpContent(title: "Preview Count", description: description, showing: showing) +} + func seedHelp(_ showing: Binding) -> some View { let description = """ @@ -119,7 +138,7 @@ func seedHelp(_ showing: Binding) -> some View { Next, maybe add a negative prompt or tweak the prompt slightly, and see how the result changes. \ Rinse and repeat until you are satisfied, or select a new seed to start over. - If you select -1, a random seed will be chosen for you. + Set the value to 0 for a random seed to be chosen for you. """ return helpContent(title: "Generation Seed", description: description, showing: showing) } diff --git a/Diffusion-macOS/StatusView.swift b/Diffusion-macOS/StatusView.swift index ac02f42..e255dba 100644 --- a/Diffusion-macOS/StatusView.swift +++ b/Diffusion-macOS/StatusView.swift @@ -87,10 +87,11 @@ struct StatusView: View { let intervalString = String(format: "Time: %.1fs", interval ?? 0) Text(intervalString) Spacer() - if generation.seed != Double(lastSeed) { - Text("Seed: \(lastSeed)") + if generation.seed != lastSeed { + + Text(String("Seed: \(formatLargeNumber(lastSeed))")) Button("Set") { - generation.seed = Double(lastSeed) + generation.seed = lastSeed } } }.frame(maxHeight: 25) diff --git a/Diffusion-macOS/Utils_macOS.swift b/Diffusion-macOS/Utils_macOS.swift new file mode 100644 index 0000000..ea91585 --- /dev/null +++ b/Diffusion-macOS/Utils_macOS.swift @@ -0,0 +1,17 @@ +// +// Utils_macOS.swift +// Diffusion-macOS +// +// Created by Dolmere on 31/07/2023. +// + +import SwiftUI + +extension CGImage { + static func fromData(_ imageData: Data) -> CGImage? { + if let image = NSBitmapImageRep(data: imageData)?.cgImage { + return image + } + return nil + } +} diff --git a/Diffusion.xcodeproj/project.pbxproj b/Diffusion.xcodeproj/project.pbxproj index 68fe22d..dccf3de 100644 --- a/Diffusion.xcodeproj/project.pbxproj +++ b/Diffusion.xcodeproj/project.pbxproj @@ -7,18 +7,24 @@ objects = { /* Begin PBXBuildFile section */ + 16AFDD4F2C1B7D6200536A62 /* StableDiffusion in Frameworks */ = {isa = PBXBuildFile; productRef = 16AFDD4E2C1B7D6200536A62 /* StableDiffusion */; }; + 16AFDD512C1B7D6700536A62 /* StableDiffusion in Frameworks */ = {isa = PBXBuildFile; productRef = 16AFDD502C1B7D6700536A62 /* StableDiffusion */; }; + 8C4B32042A770C1D0090EF17 /* DiffusionImage+macOS.swift in Sources */ = {isa = PBXBuildFile; fileRef = 8C4B32032A770C1D0090EF17 /* DiffusionImage+macOS.swift */; }; + 8C4B32062A770C300090EF17 /* DiffusionImage+iOS.swift in Sources */ = {isa = PBXBuildFile; fileRef = 8C4B32052A770C300090EF17 /* DiffusionImage+iOS.swift */; }; + 8C4B32082A77F90C0090EF17 /* Utils_iOS.swift in Sources */ = {isa = PBXBuildFile; fileRef = 8C4B32072A77F90C0090EF17 /* Utils_iOS.swift */; }; + 8C4B320A2A77F9160090EF17 /* Utils_macOS.swift in Sources */ = {isa = PBXBuildFile; fileRef = 8C4B32092A77F9160090EF17 /* Utils_macOS.swift */; }; + 8CD8A53A2A456EF800BD8A98 /* PromptTextField.swift in Sources */ = {isa = PBXBuildFile; fileRef = 8CD8A5392A456EF800BD8A98 /* PromptTextField.swift */; }; + 8CD8A53C2A476E2C00BD8A98 /* PromptTextField.swift in Sources */ = {isa = PBXBuildFile; fileRef = 8CD8A5392A456EF800BD8A98 /* PromptTextField.swift */; }; + 8CEEB7D92A54C88C00C23829 /* DiffusionImage.swift in Sources */ = {isa = PBXBuildFile; fileRef = 8CEEB7D82A54C88C00C23829 /* DiffusionImage.swift */; }; + 8CEEB7DA2A54C88C00C23829 /* DiffusionImage.swift in Sources */ = {isa = PBXBuildFile; fileRef = 8CEEB7D82A54C88C00C23829 /* DiffusionImage.swift */; }; EB067F872992E561004D1AD9 /* HelpContent.swift in Sources */ = {isa = PBXBuildFile; fileRef = EB067F862992E561004D1AD9 /* HelpContent.swift */; }; - EB33A51D2954D89F00B16357 /* StableDiffusion in Frameworks */ = {isa = PBXBuildFile; productRef = EB33A51C2954D89F00B16357 /* StableDiffusion */; }; EB560F0429A3C20800C0F8B8 /* Capabilities.swift in Sources */ = {isa = PBXBuildFile; fileRef = EB560F0329A3C20800C0F8B8 /* Capabilities.swift */; }; EBB5BA5329425BEE003A2A5B /* PipelineLoader.swift in Sources */ = {isa = PBXBuildFile; fileRef = EBB5BA5229425BEE003A2A5B /* PipelineLoader.swift */; }; - 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 */; }; EBDD7DAA29731F6C00C1C4B2 /* Pipeline.swift in Sources */ = {isa = PBXBuildFile; fileRef = EBE75601293E91E200806B32 /* Pipeline.swift */; }; EBDD7DAB29731F7500C1C4B2 /* PipelineLoader.swift in Sources */ = {isa = PBXBuildFile; fileRef = EBB5BA5229425BEE003A2A5B /* PipelineLoader.swift */; }; - EBDD7DAD29731FB300C1C4B2 /* Path in Frameworks */ = {isa = PBXBuildFile; productRef = EBDD7DAC29731FB300C1C4B2 /* Path */; }; EBDD7DAF29731FB300C1C4B2 /* ZIPFoundation in Frameworks */ = {isa = PBXBuildFile; productRef = EBDD7DAE29731FB300C1C4B2 /* ZIPFoundation */; }; - EBDD7DB129731FB300C1C4B2 /* StableDiffusion in Frameworks */ = {isa = PBXBuildFile; productRef = EBDD7DB029731FB300C1C4B2 /* StableDiffusion */; }; EBDD7DB32973200200C1C4B2 /* Utils.swift in Sources */ = {isa = PBXBuildFile; fileRef = EBDD7DB22973200200C1C4B2 /* Utils.swift */; }; EBDD7DB42973200200C1C4B2 /* Utils.swift in Sources */ = {isa = PBXBuildFile; fileRef = EBDD7DB22973200200C1C4B2 /* Utils.swift */; }; EBDD7DB52973201800C1C4B2 /* ModelInfo.swift in Sources */ = {isa = PBXBuildFile; fileRef = EBE3FF4B295E1EFE00E921AA /* ModelInfo.swift */; }; @@ -63,6 +69,12 @@ /* End PBXContainerItemProxy section */ /* Begin PBXFileReference section */ + 8C4B32032A770C1D0090EF17 /* DiffusionImage+macOS.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = "DiffusionImage+macOS.swift"; sourceTree = ""; }; + 8C4B32052A770C300090EF17 /* DiffusionImage+iOS.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = "DiffusionImage+iOS.swift"; sourceTree = ""; }; + 8C4B32072A77F90C0090EF17 /* Utils_iOS.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = Utils_iOS.swift; sourceTree = ""; }; + 8C4B32092A77F9160090EF17 /* Utils_macOS.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = Utils_macOS.swift; sourceTree = ""; }; + 8CD8A5392A456EF800BD8A98 /* PromptTextField.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = PromptTextField.swift; sourceTree = ""; }; + 8CEEB7D82A54C88C00C23829 /* DiffusionImage.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = DiffusionImage.swift; sourceTree = ""; }; EB067F862992E561004D1AD9 /* HelpContent.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = HelpContent.swift; sourceTree = ""; }; EB33A51E2954E1BC00B16357 /* Info.plist */ = {isa = PBXFileReference; lastKnownFileType = text.plist; path = Info.plist; sourceTree = ""; }; EB560F0329A3C20800C0F8B8 /* Capabilities.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = Capabilities.swift; sourceTree = ""; }; @@ -104,9 +116,8 @@ isa = PBXFrameworksBuildPhase; buildActionMask = 2147483647; files = ( - EBB5BA5829425E17003A2A5B /* Path in Frameworks */, - EB33A51D2954D89F00B16357 /* StableDiffusion in Frameworks */, EBB5BA5D294504DE003A2A5B /* ZIPFoundation in Frameworks */, + 16AFDD512C1B7D6700536A62 /* StableDiffusion in Frameworks */, ); runOnlyForDeploymentPostprocessing = 0; }; @@ -129,8 +140,7 @@ buildActionMask = 2147483647; files = ( F155203C297118E700DC009B /* CompactSlider in Frameworks */, - EBDD7DB129731FB300C1C4B2 /* StableDiffusion in Frameworks */, - EBDD7DAD29731FB300C1C4B2 /* Path in Frameworks */, + 16AFDD4F2C1B7D6200536A62 /* StableDiffusion in Frameworks */, EBDD7DAF29731FB300C1C4B2 /* ZIPFoundation in Frameworks */, ); runOnlyForDeploymentPostprocessing = 0; @@ -138,6 +148,29 @@ /* End PBXFrameworksBuildPhase section */ /* Begin PBXGroup section */ + 8CD8A53B2A476E1C00BD8A98 /* Views */ = { + isa = PBXGroup; + children = ( + 8CD8A5392A456EF800BD8A98 /* PromptTextField.swift */, + ); + path = Views; + sourceTree = ""; + }; + 8CF53E022A44AE0400E6358B /* Common */ = { + isa = PBXGroup; + children = ( + EBB5BA5929426E06003A2A5B /* Downloader.swift */, + EBE3FF4B295E1EFE00E921AA /* ModelInfo.swift */, + EBDD7DB72976AAFE00C1C4B2 /* State.swift */, + EBDD7DB22973200200C1C4B2 /* Utils.swift */, + 8CEEB7D82A54C88C00C23829 /* DiffusionImage.swift */, + EBB5BA5129425B07003A2A5B /* Pipeline */, + 8CD8A53B2A476E1C00BD8A98 /* Views */, + ); + name = Common; + path = Diffusion/Common; + sourceTree = ""; + }; EBB5BA5129425B07003A2A5B /* Pipeline */ = { isa = PBXGroup; children = ( @@ -155,6 +188,7 @@ EBE4438729488DCA00CDA605 /* README.md */, EBE443892948953600CDA605 /* LICENSE */, EBE755FF293E910800806B32 /* Packages */, + 8CF53E022A44AE0400E6358B /* Common */, EBE755C7293E37DD00806B32 /* Diffusion */, F15520222971093300DC009B /* Diffusion-macOS */, EBE755D9293E37DE00806B32 /* DiffusionTests */, @@ -179,16 +213,13 @@ isa = PBXGroup; children = ( EB33A51E2954E1BC00B16357 /* Info.plist */, - EBB5BA5129425B07003A2A5B /* Pipeline */, EBE7560A29411A5E00806B32 /* Views */, - EBB5BA5929426E06003A2A5B /* Downloader.swift */, EBE755C8293E37DD00806B32 /* DiffusionApp.swift */, - EBDD7DB22973200200C1C4B2 /* Utils.swift */, - EBE3FF4B295E1EFE00E921AA /* ModelInfo.swift */, - EBDD7DB72976AAFE00C1C4B2 /* State.swift */, EBE755CC293E37DD00806B32 /* Assets.xcassets */, EBE755CE293E37DD00806B32 /* Diffusion.entitlements */, EBE755CF293E37DD00806B32 /* Preview Content */, + 8C4B32052A770C300090EF17 /* DiffusionImage+iOS.swift */, + 8C4B32072A77F90C0090EF17 /* Utils_iOS.swift */, ); path = Diffusion; sourceTree = ""; @@ -252,6 +283,8 @@ F155203329710B3600DC009B /* StatusView.swift */, EB067F862992E561004D1AD9 /* HelpContent.swift */, EB560F0329A3C20800C0F8B8 /* Capabilities.swift */, + 8C4B32032A770C1D0090EF17 /* DiffusionImage+macOS.swift */, + 8C4B32092A77F9160090EF17 /* Utils_macOS.swift */, F155202C2971093400DC009B /* Diffusion_macOS.entitlements */, F15520292971093400DC009B /* Preview Content */, ); @@ -280,12 +313,12 @@ buildRules = ( ); dependencies = ( + EBF61AB32A2F976600482CF3 /* PBXTargetDependency */, ); name = Diffusion; packageProductDependencies = ( - EBB5BA5729425E17003A2A5B /* Path */, EBB5BA5C294504DE003A2A5B /* ZIPFoundation */, - EB33A51C2954D89F00B16357 /* StableDiffusion */, + 16AFDD502C1B7D6700536A62 /* StableDiffusion */, ); productName = Diffusion; productReference = EBE755C5293E37DD00806B32 /* Diffusion.app */; @@ -339,13 +372,13 @@ buildRules = ( ); dependencies = ( + EB0199492A31FEAF00B133E2 /* PBXTargetDependency */, ); name = "Diffusion-macOS"; packageProductDependencies = ( F155203B297118E700DC009B /* CompactSlider */, - EBDD7DAC29731FB300C1C4B2 /* Path */, EBDD7DAE29731FB300C1C4B2 /* ZIPFoundation */, - EBDD7DB029731FB300C1C4B2 /* StableDiffusion */, + 16AFDD4E2C1B7D6200536A62 /* StableDiffusion */, ); productName = "Diffusion-macOS"; productReference = F15520212971093300DC009B /* Diffusers.app */; @@ -387,10 +420,9 @@ ); mainGroup = EBE755BC293E37DD00806B32; packageReferences = ( - EBB5BA5629425E17003A2A5B /* XCRemoteSwiftPackageReference "Path.swift" */, EBB5BA5B294504DE003A2A5B /* XCRemoteSwiftPackageReference "ZIPFoundation" */, - EB33A51B2954D89F00B16357 /* XCRemoteSwiftPackageReference "ml-stable-diffusion" */, F155203A297118E600DC009B /* XCRemoteSwiftPackageReference "CompactSlider" */, + 16AFDD4D2C1B7D4800536A62 /* XCRemoteSwiftPackageReference "ml-stable-diffusion" */, ); productRefGroup = EBE755C6293E37DD00806B32 /* Products */; projectDirPath = ""; @@ -468,10 +500,14 @@ EBE75602293E91E200806B32 /* Pipeline.swift in Sources */, EBE755CB293E37DD00806B32 /* TextToImage.swift in Sources */, EBB5BA5A29426E06003A2A5B /* Downloader.swift in Sources */, + 8C4B32062A770C300090EF17 /* DiffusionImage+iOS.swift in Sources */, + 8CEEB7D92A54C88C00C23829 /* DiffusionImage.swift in Sources */, EBE3FF4C295E1EFE00E921AA /* ModelInfo.swift in Sources */, EBE756092941178600806B32 /* Loading.swift in Sources */, + 8C4B32082A77F90C0090EF17 /* Utils_iOS.swift in Sources */, EBDD7DB82976AAFE00C1C4B2 /* State.swift in Sources */, EBB5BA5329425BEE003A2A5B /* PipelineLoader.swift in Sources */, + 8CD8A53C2A476E2C00BD8A98 /* PromptTextField.swift in Sources */, EBE755C9293E37DD00806B32 /* DiffusionApp.swift in Sources */, EBDD7DB32973200200C1C4B2 /* Utils.swift in Sources */, ); @@ -503,20 +539,28 @@ F15520262971093300DC009B /* ContentView.swift in Sources */, EBDD7DB92976AAFE00C1C4B2 /* State.swift in Sources */, EB067F872992E561004D1AD9 /* HelpContent.swift in Sources */, + 8C4B320A2A77F9160090EF17 /* Utils_macOS.swift in Sources */, EBDD7DB42973200200C1C4B2 /* Utils.swift in Sources */, + 8CD8A53A2A456EF800BD8A98 /* PromptTextField.swift in Sources */, F1552031297109C300DC009B /* ControlsView.swift in Sources */, EBDD7DB62973206600C1C4B2 /* Downloader.swift in Sources */, F155203429710B3600DC009B /* StatusView.swift in Sources */, EB560F0429A3C20800C0F8B8 /* Capabilities.swift in Sources */, F15520242971093300DC009B /* Diffusion_macOSApp.swift in Sources */, EBDD7DB52973201800C1C4B2 /* ModelInfo.swift in Sources */, + 8C4B32042A770C1D0090EF17 /* DiffusionImage+macOS.swift in Sources */, EBDD7DBD2977FFB300C1C4B2 /* GeneratedImageView.swift in Sources */, + 8CEEB7DA2A54C88C00C23829 /* DiffusionImage.swift in Sources */, ); runOnlyForDeploymentPostprocessing = 0; }; /* End PBXSourcesBuildPhase section */ /* Begin PBXTargetDependency section */ + EB0199492A31FEAF00B133E2 /* PBXTargetDependency */ = { + isa = PBXTargetDependency; + productRef = EB0199482A31FEAF00B133E2 /* StableDiffusion */; + }; EBE755D8293E37DE00806B32 /* PBXTargetDependency */ = { isa = PBXTargetDependency; target = EBE755C4293E37DD00806B32 /* Diffusion */; @@ -527,6 +571,10 @@ target = EBE755C4293E37DD00806B32 /* Diffusion */; targetProxy = EBE755E1293E37DE00806B32 /* PBXContainerItemProxy */; }; + EBF61AB32A2F976600482CF3 /* PBXTargetDependency */ = { + isa = PBXTargetDependency; + productRef = EBF61AB22A2F976600482CF3 /* StableDiffusion */; + }; /* End PBXTargetDependency section */ /* Begin XCBuildConfiguration section */ @@ -828,7 +876,7 @@ "$(inherited)", "@executable_path/../Frameworks", ); - MACOSX_DEPLOYMENT_TARGET = 13.1; + MACOSX_DEPLOYMENT_TARGET = 14.0; SDKROOT = macosx; SWIFT_EMIT_LOC_STRINGS = YES; SWIFT_VERSION = 5.0; @@ -856,7 +904,7 @@ "$(inherited)", "@executable_path/../Frameworks", ); - MACOSX_DEPLOYMENT_TARGET = 13.1; + MACOSX_DEPLOYMENT_TARGET = 14.0; PRODUCT_BUNDLE_IDENTIFIER = com.huggingface.Diffusers; SDKROOT = macosx; SWIFT_EMIT_LOC_STRINGS = YES; @@ -915,22 +963,14 @@ /* End XCConfigurationList section */ /* Begin XCRemoteSwiftPackageReference section */ - EB33A51B2954D89F00B16357 /* XCRemoteSwiftPackageReference "ml-stable-diffusion" */ = { + 16AFDD4D2C1B7D4800536A62 /* XCRemoteSwiftPackageReference "ml-stable-diffusion" */ = { isa = XCRemoteSwiftPackageReference; - repositoryURL = "https://github.com/apple/ml-stable-diffusion"; + repositoryURL = "https://github.com/argmaxinc/ml-stable-diffusion.git"; requirement = { branch = main; kind = branch; }; }; - EBB5BA5629425E17003A2A5B /* XCRemoteSwiftPackageReference "Path.swift" */ = { - isa = XCRemoteSwiftPackageReference; - repositoryURL = "https://github.com/mxcl/Path.swift.git"; - requirement = { - kind = upToNextMajorVersion; - minimumVersion = 1.0.0; - }; - }; EBB5BA5B294504DE003A2A5B /* XCRemoteSwiftPackageReference "ZIPFoundation" */ = { isa = XCRemoteSwiftPackageReference; repositoryURL = "https://github.com/weichsel/ZIPFoundation.git"; @@ -950,34 +990,32 @@ /* End XCRemoteSwiftPackageReference section */ /* Begin XCSwiftPackageProductDependency section */ - EB33A51C2954D89F00B16357 /* StableDiffusion */ = { + 16AFDD4E2C1B7D6200536A62 /* StableDiffusion */ = { isa = XCSwiftPackageProductDependency; - package = EB33A51B2954D89F00B16357 /* XCRemoteSwiftPackageReference "ml-stable-diffusion" */; + package = 16AFDD4D2C1B7D4800536A62 /* XCRemoteSwiftPackageReference "ml-stable-diffusion" */; productName = StableDiffusion; }; - EBB5BA5729425E17003A2A5B /* Path */ = { + 16AFDD502C1B7D6700536A62 /* StableDiffusion */ = { isa = XCSwiftPackageProductDependency; - package = EBB5BA5629425E17003A2A5B /* XCRemoteSwiftPackageReference "Path.swift" */; - productName = Path; + package = 16AFDD4D2C1B7D4800536A62 /* XCRemoteSwiftPackageReference "ml-stable-diffusion" */; + productName = StableDiffusion; + }; + EB0199482A31FEAF00B133E2 /* StableDiffusion */ = { + isa = XCSwiftPackageProductDependency; + productName = StableDiffusion; }; EBB5BA5C294504DE003A2A5B /* ZIPFoundation */ = { isa = XCSwiftPackageProductDependency; package = EBB5BA5B294504DE003A2A5B /* XCRemoteSwiftPackageReference "ZIPFoundation" */; productName = ZIPFoundation; }; - EBDD7DAC29731FB300C1C4B2 /* Path */ = { - isa = XCSwiftPackageProductDependency; - package = EBB5BA5629425E17003A2A5B /* XCRemoteSwiftPackageReference "Path.swift" */; - productName = Path; - }; EBDD7DAE29731FB300C1C4B2 /* ZIPFoundation */ = { isa = XCSwiftPackageProductDependency; package = EBB5BA5B294504DE003A2A5B /* XCRemoteSwiftPackageReference "ZIPFoundation" */; productName = ZIPFoundation; }; - EBDD7DB029731FB300C1C4B2 /* StableDiffusion */ = { + EBF61AB22A2F976600482CF3 /* StableDiffusion */ = { isa = XCSwiftPackageProductDependency; - package = EB33A51B2954D89F00B16357 /* XCRemoteSwiftPackageReference "ml-stable-diffusion" */; productName = StableDiffusion; }; F155203B297118E700DC009B /* CompactSlider */ = { diff --git a/Diffusion.xcodeproj/project.xcworkspace/xcshareddata/swiftpm/Package.resolved b/Diffusion.xcodeproj/project.xcworkspace/xcshareddata/swiftpm/Package.resolved index 6f69852..3eaff71 100644 --- a/Diffusion.xcodeproj/project.xcworkspace/xcshareddata/swiftpm/Package.resolved +++ b/Diffusion.xcodeproj/project.xcworkspace/xcshareddata/swiftpm/Package.resolved @@ -1,4 +1,5 @@ { + "originHash" : "e97aab54879429ea40e58df49ffe4eef5228d95a28a7cf4d5dca9204c33564e1", "pins" : [ { "identity" : "compactslider", @@ -6,25 +7,16 @@ "location" : "https://github.com/buh/CompactSlider.git", "state" : { "branch" : "main", - "revision" : "3cb37fb7385913835b6844c6af2680c64000dcd2" + "revision" : "31a4db8ef10f32e574be35399b10b00e05d27e38" } }, { "identity" : "ml-stable-diffusion", "kind" : "remoteSourceControl", - "location" : "https://github.com/apple/ml-stable-diffusion", + "location" : "https://github.com/argmaxinc/ml-stable-diffusion.git", "state" : { "branch" : "main", - "revision" : "fb1fa01c9d30e9b2e02a8b7ed35d905e272a0262" - } - }, - { - "identity" : "path.swift", - "kind" : "remoteSourceControl", - "location" : "https://github.com/mxcl/Path.swift.git", - "state" : { - "revision" : "9c6f807b0a76be0e27aecc908bc6f173400d839e", - "version" : "1.4.0" + "revision" : "d1f0604fab5345011e0b9f5b87ee0c155612565f" } }, { @@ -32,8 +24,8 @@ "kind" : "remoteSourceControl", "location" : "https://github.com/apple/swift-argument-parser.git", "state" : { - "revision" : "fddd1c00396eed152c45a46bea9f47b98e59301d", - "version" : "1.2.0" + "revision" : "0fbc8848e389af3bb55c182bc19ca9d5dc2f255b", + "version" : "1.4.0" } }, { @@ -41,10 +33,10 @@ "kind" : "remoteSourceControl", "location" : "https://github.com/weichsel/ZIPFoundation.git", "state" : { - "revision" : "43ec568034b3731101dbf7670765d671c30f54f3", - "version" : "0.9.16" + "revision" : "b979e8b52c7ae7f3f39fa0182e738e9e7257eb78", + "version" : "0.9.18" } } ], - "version" : 2 + "version" : 3 } diff --git a/Diffusion/Assets.xcassets/placeholder.imageset/labrador.png b/Diffusion/Assets.xcassets/placeholder.imageset/labrador.png index 0cdaef7..a93e1bf 100644 Binary files a/Diffusion/Assets.xcassets/placeholder.imageset/labrador.png and b/Diffusion/Assets.xcassets/placeholder.imageset/labrador.png differ diff --git a/Diffusion/Common/DiffusionImage.swift b/Diffusion/Common/DiffusionImage.swift new file mode 100644 index 0000000..9522b84 --- /dev/null +++ b/Diffusion/Common/DiffusionImage.swift @@ -0,0 +1,131 @@ +// +// DiffusionImage.swift +// Diffusion +// +// Created by Dolmere on 03/07/2023. +// + +import SwiftUI +import StableDiffusion +import CoreTransferable + +/// Tracking for a `DiffusionImage` generation state. +enum DiffusionImageState { + case generating + case waiting + case complete +} + +/// Generic custom error to use when an image generation fails. +enum DiffusionImageError: Error { + case invalidDiffusionImage +} + +/// Combination of a `DiffusionImage` and its associated `DiffusionImageState` +struct DiffusionImageWrapper { + var diffusionImageState: DiffusionImageState = .waiting + var diffusionImage: DiffusionImage? = nil +} + +/// Model class to hold a generated image and the "recipe" data that was used to generate it +final class DiffusionImage: NSObject, Identifiable, NSCoding, NSSecureCoding { + + let id: UUID + let cgImage: CGImage + let seed: UInt32 + let steps: Double + let positivePrompt: String + let negativePrompt: String + let guidanceScale: Double + let disableSafety: Bool + /// Local enum represented with a String to conform to NSSecureCoding + let scheduler: StableDiffusionScheduler + + /// This is a composed `String` built from the numeric `Seed` and the user supplied `positivePrompt` limited to the first 200 character and with whitespace replaced with underscore characters. + var generatedFilename: String { + return "\(seed)-\(positivePrompt)".first200Safe + } + + /// The location on the file system where this generated image is stored. + var fileURL: URL + + init(id: UUID, cgImage: CGImage, seed: UInt32, steps: Double, positivePrompt: String, negativePrompt: String, guidanceScale: Double, disableSafety: Bool, scheduler: StableDiffusionScheduler) { + let genname = "\(seed)-\(positivePrompt)".first200Safe + self.id = id + self.cgImage = cgImage + self.seed = seed + self.steps = steps + self.positivePrompt = positivePrompt + self.negativePrompt = negativePrompt + self.guidanceScale = guidanceScale + self.disableSafety = disableSafety + self.scheduler = scheduler + // Initially set the fileURL to the top level applicationDirectory to allow running the completed instance func save() where the fileURL will be updated to the correct location. + self.fileURL = URL.applicationDirectory + // init the instance fully before executing an instance function + super.init() + if let url = save(cgImage: cgImage, filename: genname) { + self.fileURL = url + } else { + fatalError("Fatal error init of DiffusionImage, cannot create image file at \(genname)") + } + } + + func encode(with coder: NSCoder) { + coder.encode(id, forKey: "id") + coder.encode(seed, forKey: "seed") + coder.encode(steps, forKey: "steps") + coder.encode(positivePrompt, forKey: "positivePrompt") + coder.encode(negativePrompt, forKey: "negativePrompt") + coder.encode(guidanceScale, forKey: "guidanceScale") + coder.encode(disableSafety, forKey: "disableSafety") + coder.encode(scheduler, forKey: "scheduler") + // Encode cgImage as data + if let data = pngRepresentation() { + coder.encode(data, forKey: "cgImage") + } + } + + required init?(coder: NSCoder) { + guard let id = coder.decodeObject(forKey: "id") as? UUID else { + return nil + } + + self.id = id + self.seed = UInt32(coder.decodeInt32(forKey: "seed")) + self.steps = coder.decodeDouble(forKey: "steps") + self.positivePrompt = coder.decodeObject(forKey: "positivePrompt") as? String ?? "" + self.negativePrompt = coder.decodeObject(forKey: "negativePrompt") as? String ?? "" + self.guidanceScale = coder.decodeDouble(forKey: "guidanceScale") + self.disableSafety = coder.decodeBool(forKey: "disableSafety") + self.scheduler = coder.decodeObject(forKey: "scheduler") as? StableDiffusionScheduler ?? StableDiffusionScheduler.dpmSolverMultistepScheduler + let genname = "\(seed)-\(positivePrompt)".first200Safe + + // Decode cgImage from data + if let imageData = coder.decodeObject(forKey: "cgImage") as? Data { + guard let img = CGImage.fromData(imageData) else { fatalError("Fatal error loading data with missing cgImage in object") } + self.cgImage = img + } else { + fatalError("Fatal error loading data with missing cgImage in object") + } + self.fileURL = URL.applicationDirectory + super.init() + if let url = save(cgImage: cgImage, filename: genname) { + self.fileURL = url + } else { + fatalError("Fatal error init of DiffusionImage, cannot create image file at \(genname)") + } + } + + // MARK: - Equatable + + static func == (lhs: DiffusionImage, rhs: DiffusionImage) -> Bool { + return lhs.id == rhs.id + } + + // MARK: - NSSecureCoding + + static var supportsSecureCoding: Bool { + return true + } +} diff --git a/Diffusion/Downloader.swift b/Diffusion/Common/Downloader.swift similarity index 65% rename from Diffusion/Downloader.swift rename to Diffusion/Common/Downloader.swift index 8bb8e4c..b8cd9cd 100644 --- a/Diffusion/Downloader.swift +++ b/Diffusion/Common/Downloader.swift @@ -8,7 +8,6 @@ import Foundation import Combine -import Path class Downloader: NSObject, ObservableObject { private(set) var destination: URL @@ -25,12 +24,18 @@ class Downloader: NSObject, ObservableObject { private var urlSession: URLSession? = nil - init(from url: URL, to destination: URL) { + init(from url: URL, to destination: URL, using authToken: String? = nil) { self.destination = destination super.init() + var config = URLSessionConfiguration.default + #if !os(macOS) // .background allows downloads to proceed in the background - let config = URLSessionConfiguration.background(withIdentifier: "net.pcuenca.diffusion.download") + // helpful for devices that may not keep the app in the foreground for the download duration + config = URLSessionConfiguration.background(withIdentifier: "net.pcuenca.diffusion.download") + config.isDiscretionary = false + config.sessionSendsLaunchEvents = true + #endif urlSession = URLSession(configuration: config, delegate: self, delegateQueue: OperationQueue()) downloadState.value = .downloading(0) urlSession?.getAllTasks { tasks in @@ -40,7 +45,13 @@ class Downloader: NSObject, ObservableObject { return } print("Starting download of \(url)") - self.urlSession?.downloadTask(with: url).resume() + + var request = URLRequest(url: url) + if let authToken = authToken { + request.setValue("Bearer \(authToken)", forHTTPHeaderField: "Authorization") + } + + self.urlSession?.downloadTask(with: request).resume() } } @@ -70,30 +81,30 @@ class Downloader: NSObject, ObservableObject { } extension Downloader: URLSessionDelegate, URLSessionDownloadDelegate { - func urlSession(_: URLSession, downloadTask: URLSessionDownloadTask, didWriteData _: Int64, totalBytesWritten _: Int64, totalBytesExpectedToWrite _: Int64) { - downloadState.value = .downloading(downloadTask.progress.fractionCompleted) + func urlSession(_: URLSession, downloadTask: URLSessionDownloadTask, didWriteData _: Int64, totalBytesWritten: Int64, totalBytesExpectedToWrite: Int64) { + downloadState.value = .downloading(Double(totalBytesWritten) / Double(totalBytesExpectedToWrite)) } func urlSession(_: URLSession, downloadTask _: URLSessionDownloadTask, didFinishDownloadingTo location: URL) { - guard let path = Path(url: location) else { + guard FileManager.default.fileExists(atPath: location.path) else { downloadState.value = .failed("Invalid download location received: \(location)") return } - guard let toPath = Path(url: destination) else { - downloadState.value = .failed("Invalid destination: \(destination)") - return - } do { - try path.move(to: toPath, overwrite: true) + try FileManager.default.moveItem(at: location, to: destination) downloadState.value = .completed(destination) } catch { downloadState.value = .failed(error) } } - func urlSession(_: URLSession, task: URLSessionTask, didCompleteWithError error: Error?) { + func urlSession(_ session: URLSession, task: URLSessionTask, didCompleteWithError error: Error?) { if let error = error { downloadState.value = .failed(error) + } else if let response = task.response as? HTTPURLResponse { + print("HTTP response status code: \(response.statusCode)") +// let headers = response.allHeaderFields +// print("HTTP response headers: \(headers)") } } } diff --git a/Diffusion/Common/ModelInfo.swift b/Diffusion/Common/ModelInfo.swift new file mode 100644 index 0000000..bffcd69 --- /dev/null +++ b/Diffusion/Common/ModelInfo.swift @@ -0,0 +1,279 @@ +// +// ModelInfo.swift +// Diffusion +// +// Created by Pedro Cuenca on 29/12/22. +// See LICENSE at https://github.com/huggingface/swift-coreml-diffusers/LICENSE +// + +import CoreML + +enum AttentionVariant: String { + case original + case splitEinsum + case splitEinsumV2 +} + +extension AttentionVariant { + var defaultComputeUnits: MLComputeUnits { self == .original ? .cpuAndGPU : .cpuAndNeuralEngine } +} + +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 splitAttentionSuffix: String + + /// Suffix of the archive containing the SPLIT_EINSUM_V2 attention variant. Usually something like "split_einsum_v2_compiled" + let splitAttentionV2Suffix: String + + /// Whether the archive contains ANE optimized models + let supportsNeuralEngine: Bool + + /// Whether the archive contains the VAE Encoder (for image to image tasks). Not yet in use. + let supportsEncoder: Bool + + /// Is attention v2 supported? (Ideally, we should know by looking at the repo contents) + let supportsAttentionV2: Bool + + /// Are weights quantized? This is only used to decide whether to use `reduceMemory` + let quantized: Bool + + /// Whether this is a Stable Diffusion XL model + // TODO: retrieve from remote config + let isXL: Bool + + /// Whether this is a Stable Diffusion 3 model + // TODO: retrieve from remote config + let isSD3: Bool + + //TODO: refactor all these properties + init(modelId: String, modelVersion: String, + originalAttentionSuffix: String = "original_compiled", + splitAttentionSuffix: String = "split_einsum_compiled", + splitAttentionV2Suffix: String = "split_einsum_v2_compiled", + supportsNeuralEngine: Bool = true, + supportsEncoder: Bool = false, + supportsAttentionV2: Bool = false, + quantized: Bool = false, + isXL: Bool = false, + isSD3: Bool = false) { + self.modelId = modelId + self.modelVersion = modelVersion + self.originalAttentionSuffix = originalAttentionSuffix + self.splitAttentionSuffix = splitAttentionSuffix + self.splitAttentionV2Suffix = splitAttentionV2Suffix + self.supportsNeuralEngine = supportsNeuralEngine + self.supportsEncoder = supportsEncoder + self.supportsAttentionV2 = supportsAttentionV2 + self.quantized = quantized + self.isXL = isXL + self.isSD3 = isSD3 + } +} + +extension ModelInfo { + //TODO: set compute units instead and derive variant from it + static var defaultAttention: AttentionVariant { + guard runningOnMac else { return .splitEinsum } + #if os(macOS) + guard Capabilities.hasANE else { return .original } + return Capabilities.performanceCores >= 8 ? .original : .splitEinsum + #else + return .splitEinsum + #endif + } + + static var defaultComputeUnits: MLComputeUnits { defaultAttention.defaultComputeUnits } + + var bestAttention: AttentionVariant { + if !runningOnMac && supportsAttentionV2 { return .splitEinsumV2 } + return ModelInfo.defaultAttention + } + var defaultComputeUnits: MLComputeUnits { bestAttention.defaultComputeUnits } + + 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 + switch variant { + case .original: suffix = originalAttentionSuffix + case .splitEinsum: suffix = splitAttentionSuffix + case .splitEinsumV2: suffix = splitAttentionV2Suffix + } + let repo = modelId.split(separator: "/").last! + return URL(string: "https://huggingface.co/\(modelId)/resolve/main/\(repo)_\(suffix).zip")! + } + + /// Best variant for the current platform. + /// Currently using `split_einsum` for iOS and simple performance heuristics for macOS. + var bestURL: URL { modelURL(for: bestAttention) } + + var reduceMemory: Bool { + // Enable on iOS devices, except when using quantization + if runningOnMac { return false } + if isXL { return !deviceHas8GBOrMore } + return !(quantized && deviceHas6GBOrMore) + } +} + +extension ModelInfo { + static let v14Base = ModelInfo( + modelId: "pcuenq/coreml-stable-diffusion-1-4", + modelVersion: "CompVis SD 1.4" + ) + + static let v14Palettized = ModelInfo( + modelId: "apple/coreml-stable-diffusion-1-4-palettized", + modelVersion: "CompVis SD 1.4 [6 bit]", + supportsEncoder: true, + supportsAttentionV2: true, + quantized: true + ) + + static let v15Base = ModelInfo( + modelId: "pcuenq/coreml-stable-diffusion-v1-5", + modelVersion: "RunwayML SD 1.5" + ) + + static let v15Palettized = ModelInfo( + modelId: "apple/coreml-stable-diffusion-v1-5-palettized", + modelVersion: "RunwayML SD 1.5 [6 bit]", + supportsEncoder: true, + supportsAttentionV2: true, + quantized: true + ) + + static let v2Base = ModelInfo( + modelId: "pcuenq/coreml-stable-diffusion-2-base", + modelVersion: "StabilityAI SD 2.0", + supportsEncoder: true + ) + + static let v2Palettized = ModelInfo( + modelId: "apple/coreml-stable-diffusion-2-base-palettized", + modelVersion: "StabilityAI SD 2.0 [6 bit]", + supportsEncoder: true, + supportsAttentionV2: true, + quantized: true + ) + + static let v21Base = ModelInfo( + modelId: "pcuenq/coreml-stable-diffusion-2-1-base", + modelVersion: "StabilityAI SD 2.1", + supportsEncoder: true + ) + + static let v21Palettized = ModelInfo( + modelId: "apple/coreml-stable-diffusion-2-1-base-palettized", + modelVersion: "StabilityAI SD 2.1 [6 bit]", + supportsEncoder: true, + supportsAttentionV2: true, + quantized: true + ) + + static let ofaSmall = ModelInfo( + modelId: "pcuenq/coreml-small-stable-diffusion-v0", + modelVersion: "OFA-Sys/small-stable-diffusion-v0" + ) + + static let xl = ModelInfo( + modelId: "apple/coreml-stable-diffusion-xl-base", + modelVersion: "SDXL base (1024, macOS)", + supportsEncoder: true, + isXL: true + ) + + static let xlWithRefiner = ModelInfo( + modelId: "apple/coreml-stable-diffusion-xl-base-with-refiner", + modelVersion: "SDXL with refiner (1024, macOS)", + supportsEncoder: true, + isXL: true + ) + + static let xlmbp = ModelInfo( + modelId: "apple/coreml-stable-diffusion-mixed-bit-palettization", + modelVersion: "SDXL base (1024, macOS) [4.5 bit]", + supportsEncoder: true, + quantized: true, + isXL: true + ) + + static let xlmbpChunked = ModelInfo( + modelId: "apple/coreml-stable-diffusion-xl-base-ios", + modelVersion: "SDXL base (768, iOS) [4 bit]", + supportsEncoder: false, + quantized: true, + isXL: true + ) + + static let sd3 = ModelInfo( + modelId: "argmaxinc/coreml-stable-diffusion-3-medium", + modelVersion: "SD3 medium (512, macOS)", + supportsNeuralEngine: false, // TODO: support SD3 on ANE + supportsEncoder: false, + quantized: false, + isSD3: true + ) + + static let sd3highres = ModelInfo( + modelId: "argmaxinc/coreml-stable-diffusion-3-medium-1024-t5", + modelVersion: "SD3 medium (1024, T5, macOS)", + supportsNeuralEngine: false, // TODO: support SD3 on ANE + supportsEncoder: false, + quantized: false, + isSD3: true + ) + + static let MODELS: [ModelInfo] = { + if deviceSupportsQuantization { + var models = [ + ModelInfo.v14Base, + ModelInfo.v14Palettized, + ModelInfo.v15Base, + ModelInfo.v15Palettized, + ModelInfo.v2Base, + ModelInfo.v2Palettized, + ModelInfo.v21Base, + ModelInfo.v21Palettized + ] + if runningOnMac { + models.append(contentsOf: [ + ModelInfo.xl, + ModelInfo.xlWithRefiner, + ModelInfo.xlmbp, + ModelInfo.sd3, + ModelInfo.sd3highres, + ]) + } else { + models.append(ModelInfo.xlmbpChunked) + } + return models + } else { + return [ + ModelInfo.v14Base, + ModelInfo.v15Base, + ModelInfo.v2Base, + ModelInfo.v21Base, + ] + } + }() + + static func from(modelVersion: String) -> ModelInfo? { + ModelInfo.MODELS.first(where: {$0.modelVersion == modelVersion}) + } + + static func from(modelId: String) -> ModelInfo? { + ModelInfo.MODELS.first(where: {$0.modelId == modelId}) + } +} + +extension ModelInfo : Equatable { + static func ==(lhs: ModelInfo, rhs: ModelInfo) -> Bool { lhs.modelId == rhs.modelId } +} diff --git a/Diffusion/Common/Pipeline/Pipeline.swift b/Diffusion/Common/Pipeline/Pipeline.swift new file mode 100644 index 0000000..ae34b5b --- /dev/null +++ b/Diffusion/Common/Pipeline/Pipeline.swift @@ -0,0 +1,139 @@ +// +// Pipeline.swift +// Diffusion +// +// Created by Pedro Cuenca on December 2022. +// See LICENSE at https://github.com/huggingface/swift-coreml-diffusers/LICENSE +// + +import Foundation +import CoreML +import Combine + +import StableDiffusion + +struct StableDiffusionProgress { + var progress: StableDiffusionPipeline.Progress + + var step: Int { progress.step } + var stepCount: Int { progress.stepCount } + + var currentImages: [CGImage?] + + init(progress: StableDiffusionPipeline.Progress, previewIndices: [Bool]) { + self.progress = progress + self.currentImages = [nil] + + // Since currentImages is a computed property, only access the preview image if necessary + if progress.step < previewIndices.count, previewIndices[progress.step] { + self.currentImages = progress.currentImages + } + } +} + +struct GenerationResult { + var image: CGImage? + var lastSeed: UInt32 + var interval: TimeInterval? + var userCanceled: Bool + var itsPerSecond: Double? +} + +class Pipeline { + let pipeline: StableDiffusionPipelineProtocol + let maxSeed: UInt32 + + var isXL: Bool { + if #available(macOS 14.0, iOS 17.0, *) { + return (pipeline as? StableDiffusionXLPipeline) != nil + } + return false + } + + var isSD3: Bool { + if #available(macOS 14.0, iOS 17.0, *) { + return (pipeline as? StableDiffusion3Pipeline) != nil + } + return false + } + + var progress: StableDiffusionProgress? = nil { + didSet { + progressPublisher.value = progress + } + } + lazy private(set) var progressPublisher: CurrentValueSubject = CurrentValueSubject(progress) + + private var canceled = false + + init(_ pipeline: StableDiffusionPipelineProtocol, maxSeed: UInt32 = UInt32.max) { + self.pipeline = pipeline + self.maxSeed = maxSeed + } + + func generate( + prompt: String, + negativePrompt: String = "", + scheduler: StableDiffusionScheduler, + numInferenceSteps stepCount: Int = 50, + seed: UInt32 = 0, + numPreviews previewCount: Int = 5, + guidanceScale: Float = 7.5, + disableSafety: Bool = false + ) throws -> GenerationResult { + let beginDate = Date() + canceled = false + let theSeed = seed > 0 ? seed : UInt32.random(in: 1...maxSeed) + let sampleTimer = SampleTimer() + sampleTimer.start() + + var config = StableDiffusionPipeline.Configuration(prompt: prompt) + config.negativePrompt = negativePrompt + config.stepCount = stepCount + config.seed = theSeed + config.guidanceScale = guidanceScale + config.disableSafety = disableSafety + config.schedulerType = scheduler.asStableDiffusionScheduler() + config.useDenoisedIntermediates = true + if isXL { + config.encoderScaleFactor = 0.13025 + config.decoderScaleFactor = 0.13025 + config.schedulerTimestepSpacing = .karras + } + + if isSD3 { + config.encoderScaleFactor = 1.5305 + config.decoderScaleFactor = 1.5305 + config.decoderShiftFactor = 0.0609 + config.schedulerTimestepShift = 3.0 + } + + // Evenly distribute previews based on inference steps + let previewIndices = previewIndices(stepCount, previewCount) + + let images = try pipeline.generateImages(configuration: config) { progress in + sampleTimer.stop() + handleProgress(StableDiffusionProgress(progress: progress, + previewIndices: previewIndices), + sampleTimer: sampleTimer) + if progress.stepCount != progress.step { + sampleTimer.start() + } + return !canceled + } + let interval = Date().timeIntervalSince(beginDate) + print("Got images: \(images) in \(interval)") + + // Unwrap the 1 image we asked for, nil means safety checker triggered + let image = images.compactMap({ $0 }).first + return GenerationResult(image: image, lastSeed: theSeed, interval: interval, userCanceled: canceled, itsPerSecond: 1.0/sampleTimer.median) + } + + func handleProgress(_ progress: StableDiffusionProgress, sampleTimer: SampleTimer) { + self.progress = progress + } + + func setCancelled() { + canceled = true + } +} diff --git a/Diffusion/Pipeline/PipelineLoader.swift b/Diffusion/Common/Pipeline/PipelineLoader.swift similarity index 53% rename from Diffusion/Pipeline/PipelineLoader.swift rename to Diffusion/Common/Pipeline/PipelineLoader.swift index 34b51c2..e2d0029 100644 --- a/Diffusion/Pipeline/PipelineLoader.swift +++ b/Diffusion/Common/Pipeline/PipelineLoader.swift @@ -10,17 +10,15 @@ import CoreML import Combine -import Path import ZIPFoundation import StableDiffusion class PipelineLoader { - static let models = Path.applicationSupport / "hf-diffusion-models" - + static let models = Settings.shared.applicationSupportURL().appendingPathComponent("hf-diffusion-models") let model: ModelInfo let computeUnits: ComputeUnits let maxSeed: UInt32 - + private var downloadSubscriber: Cancellable? init(model: ModelInfo, computeUnits: ComputeUnits? = nil, maxSeed: UInt32 = UInt32.max) { @@ -30,7 +28,7 @@ class PipelineLoader { state = .undetermined setInitialState() } - + enum PipelinePreparationPhase { case undetermined case waitingToDownload @@ -41,7 +39,7 @@ class PipelineLoader { case loaded case failed(Error) } - + var state: PipelinePreparationPhase { didSet { statePublisher.value = state @@ -64,11 +62,18 @@ class PipelineLoader { } extension PipelineLoader { + // Unused. Kept for debugging purposes. --pcuenca static func removeAll() { - try? models.delete() + // Delete the parent models folder as it will be recreated when it's needed again + do { + try FileManager.default.removeItem(at: models) + } catch { + print("Failed to delete: \(models), error: \(error.localizedDescription)") + } } } + extension PipelineLoader { func cancel() { downloader?.cancel() } } @@ -77,56 +82,60 @@ extension PipelineLoader { var url: URL { return model.modelURL(for: variant) } - + var filename: String { return url.lastPathComponent } - - var downloadedPath: Path { PipelineLoader.models / filename } - var downloadedURL: URL { downloadedPath.url } - var uncompressPath: Path { downloadedPath.parent } - - var packagesFilename: String { downloadedPath.basename(dropExtension: true) } - var compiledPath: Path { downloadedPath.parent/packagesFilename } + var downloadedURL: URL { PipelineLoader.models.appendingPathComponent(filename) } + + var uncompressURL: URL { PipelineLoader.models } + + var packagesFilename: String { (filename as NSString).deletingPathExtension } + + var compiledURL: URL { downloadedURL.deletingLastPathComponent().appendingPathComponent(packagesFilename) } var downloaded: Bool { - return downloadedPath.exists + return FileManager.default.fileExists(atPath: downloadedURL.path) } - + var ready: Bool { - return compiledPath.exists + return FileManager.default.fileExists(atPath: compiledURL.path) } - + var variant: AttentionVariant { switch computeUnits { case .cpuOnly : return .original // Not supported yet case .cpuAndGPU : return .original - case .cpuAndNeuralEngine: return .splitEinsum - case .all : return .splitEinsum + case .cpuAndNeuralEngine: return model.supportsAttentionV2 ? .splitEinsumV2 : .splitEinsum + case .all : return model.isSD3 ? .original : .splitEinsum @unknown default: fatalError("Unknown MLComputeUnits") } } - - // TODO: maybe receive Progress to add another progress as child + func prepare() async throws -> Pipeline { do { - try PipelineLoader.models.mkdir(.p) + do { + try FileManager.default.createDirectory(atPath: PipelineLoader.models.path, withIntermediateDirectories: true, attributes: nil) + } catch { + print("Error creating PipelineLoader.models path: \(error)") + } + try await download() try await unzip() - let pipeline = try await load(url: compiledPath.url) + let pipeline = try await load(url: compiledURL) return Pipeline(pipeline, maxSeed: maxSeed) } catch { state = .failed(error) throw error } } - + @discardableResult func download() async throws -> URL { if ready || downloaded { return downloadedURL } - + let downloader = Downloader(from: url, to: downloadedURL) self.downloader = downloader downloadSubscriber = downloader.downloadState.sink { state in @@ -137,29 +146,51 @@ extension PipelineLoader { try downloader.waitUntilDone() return downloadedURL } - + func unzip() async throws { guard downloaded else { return } state = .uncompressing do { - try FileManager().unzipItem(at: downloadedURL, to: uncompressPath.url) + try FileManager().unzipItem(at: downloadedURL, to: uncompressURL) } catch { // Cleanup if error occurs while unzipping - try uncompressPath.delete() + try FileManager.default.removeItem(at: uncompressURL) throw error } - try downloadedPath.delete() + try FileManager.default.removeItem(at: downloadedURL) state = .readyOnDisk } - - func load(url: URL) async throws -> StableDiffusionPipeline { + + func load(url: URL) async throws -> StableDiffusionPipelineProtocol { let beginDate = Date() let configuration = MLModelConfiguration() configuration.computeUnits = computeUnits - let pipeline = try StableDiffusionPipeline(resourcesAt: url, - configuration: configuration, - disableSafety: false, - reduceMemory: model.reduceMemory) + let pipeline: StableDiffusionPipelineProtocol + if model.isXL { + if #available(macOS 14.0, iOS 17.0, *) { + pipeline = try StableDiffusionXLPipeline(resourcesAt: url, + configuration: configuration, + reduceMemory: model.reduceMemory) + } else { + throw "Stable Diffusion XL requires macOS 14" + } + + } else if model.isSD3 { + if #available(macOS 14.0, iOS 17.0, *) { + pipeline = try StableDiffusion3Pipeline(resourcesAt: url, + configuration: configuration, + reduceMemory: model.reduceMemory) + } else { + throw "Stable Diffusion 3 requires macOS 14" + } + } else { + pipeline = try StableDiffusionPipeline(resourcesAt: url, + controlNet: [], + configuration: configuration, + disableSafety: false, + reduceMemory: model.reduceMemory) + } + try pipeline.loadResources() print("Pipeline loaded in \(Date().timeIntervalSince(beginDate))") state = .loaded return pipeline diff --git a/Diffusion/Common/State.swift b/Diffusion/Common/State.swift new file mode 100644 index 0000000..ae2601a --- /dev/null +++ b/Diffusion/Common/State.swift @@ -0,0 +1,264 @@ +// +// State.swift +// Diffusion +// +// Created by Pedro Cuenca on 17/1/23. +// See LICENSE at https://github.com/huggingface/swift-coreml-diffusers/LICENSE +// + +import Combine +import SwiftUI +import StableDiffusion +import CoreML + +let DEFAULT_MODEL = ModelInfo.sd3 +let DEFAULT_PROMPT = "Labrador in the style of Vermeer" + +enum GenerationState { + case startup + case running(StableDiffusionProgress?) + case complete(String, CGImage?, UInt32, TimeInterval?) + case userCanceled + case failed(Error) +} + +typealias ComputeUnits = MLComputeUnits + +/// Schedulers compatible with StableDiffusionPipeline. This is a local implementation of the StableDiffusionScheduler enum as a String represetation to allow for compliance with NSSecureCoding. +public enum StableDiffusionScheduler: String { + /// Scheduler that uses a pseudo-linear multi-step (PLMS) method + case pndmScheduler + /// Scheduler that uses a second order DPM-Solver++ algorithm + case dpmSolverMultistepScheduler + /// Scheduler for rectified flow based multimodal diffusion transformer models + case discreteFlowScheduler + + func asStableDiffusionScheduler() -> StableDiffusion.StableDiffusionScheduler { + switch self { + case .pndmScheduler: return StableDiffusion.StableDiffusionScheduler.pndmScheduler + case .dpmSolverMultistepScheduler: return StableDiffusion.StableDiffusionScheduler.dpmSolverMultistepScheduler + case .discreteFlowScheduler: return StableDiffusion.StableDiffusionScheduler.discreteFlowScheduler + } + } +} + +class GenerationContext: ObservableObject { + let scheduler = StableDiffusionScheduler.dpmSolverMultistepScheduler + + @Published var pipeline: Pipeline? = nil { + didSet { + if let pipeline = pipeline { + progressSubscriber = pipeline + .progressPublisher + .receive(on: DispatchQueue.main) + .sink { progress in + guard let progress = progress else { return } + self.updatePreviewIfNeeded(progress) + self.state = .running(progress) + } + } + } + } + @Published var state: GenerationState = .startup + + @Published var positivePrompt = Settings.shared.prompt + @Published var negativePrompt = Settings.shared.negativePrompt + + // FIXME: Double to support the slider component + @Published var steps: Double = Settings.shared.stepCount + @Published var numImages: Double = 1.0 + @Published var seed: UInt32 = Settings.shared.seed + @Published var guidanceScale: Double = Settings.shared.guidanceScale + @Published var previews: Double = runningOnMac ? Settings.shared.previewCount : 0.0 + @Published var disableSafety = false + @Published var previewImage: CGImage? = nil + + @Published var computeUnits: ComputeUnits = Settings.shared.userSelectedComputeUnits ?? ModelInfo.defaultComputeUnits + + private var progressSubscriber: Cancellable? + + private func updatePreviewIfNeeded(_ progress: StableDiffusionProgress) { + if previews == 0 || progress.step == 0 { + previewImage = nil + } + + if previews > 0, let newImage = progress.currentImages.first, newImage != nil { + previewImage = newImage + } + } + + func generate() async throws -> GenerationResult { + guard let pipeline = pipeline else { throw "No pipeline" } + return try pipeline.generate( + prompt: positivePrompt, + negativePrompt: negativePrompt, + scheduler: scheduler, + numInferenceSteps: Int(steps), + seed: seed, + numPreviews: Int(previews), + guidanceScale: Float(guidanceScale), + disableSafety: disableSafety + ) + } + + func cancelGeneration() { + pipeline?.setCancelled() + } +} + +class Settings { + static let shared = Settings() + + let defaults = UserDefaults.standard + + enum Keys: String { + case model + case safetyCheckerDisclaimer + case computeUnits + case prompt + case negativePrompt + case guidanceScale + case stepCount + case previewCount + case seed + } + + private init() { + defaults.register(defaults: [ + Keys.model.rawValue: ModelInfo.v2Base.modelId, + Keys.safetyCheckerDisclaimer.rawValue: false, + Keys.computeUnits.rawValue: -1, // Use default + Keys.prompt.rawValue: DEFAULT_PROMPT, + Keys.negativePrompt.rawValue: "", + Keys.guidanceScale.rawValue: 7.5, + Keys.stepCount.rawValue: 25, + Keys.previewCount.rawValue: 5, + Keys.seed.rawValue: 0 + ]) + } + + var currentModel: ModelInfo { + set { + defaults.set(newValue.modelId, forKey: Keys.model.rawValue) + } + get { + guard let modelId = defaults.string(forKey: Keys.model.rawValue) else { return DEFAULT_MODEL } + return ModelInfo.from(modelId: modelId) ?? DEFAULT_MODEL + } + } + + var prompt: String { + set { + defaults.set(newValue, forKey: Keys.prompt.rawValue) + } + get { + return defaults.string(forKey: Keys.prompt.rawValue) ?? DEFAULT_PROMPT + } + } + + var negativePrompt: String { + set { + defaults.set(newValue, forKey: Keys.negativePrompt.rawValue) + } + get { + return defaults.string(forKey: Keys.negativePrompt.rawValue) ?? "" + } + } + + var guidanceScale: Double { + set { + defaults.set(newValue, forKey: Keys.guidanceScale.rawValue) + } + get { + return defaults.double(forKey: Keys.guidanceScale.rawValue) + } + } + + var stepCount: Double { + set { + defaults.set(newValue, forKey: Keys.stepCount.rawValue) + } + get { + return defaults.double(forKey: Keys.stepCount.rawValue) + } + } + + var previewCount: Double { + set { + defaults.set(newValue, forKey: Keys.previewCount.rawValue) + } + get { + return defaults.double(forKey: Keys.previewCount.rawValue) + } + } + + var seed: UInt32 { + set { + defaults.set(String(newValue), forKey: Keys.seed.rawValue) + } + get { + if let seedString = defaults.string(forKey: Keys.seed.rawValue), let seedValue = UInt32(seedString) { + return seedValue + } + return 0 + } + } + + var safetyCheckerDisclaimerShown: Bool { + set { + defaults.set(newValue, forKey: Keys.safetyCheckerDisclaimer.rawValue) + } + get { + return defaults.bool(forKey: Keys.safetyCheckerDisclaimer.rawValue) + } + } + + /// Returns the option selected by the user, if overridden + /// `nil` means: guess best + var userSelectedComputeUnits: ComputeUnits? { + set { + // Any value other than the supported ones would cause `get` to return `nil` + defaults.set(newValue?.rawValue ?? -1, forKey: Keys.computeUnits.rawValue) + } + get { + let current = defaults.integer(forKey: Keys.computeUnits.rawValue) + guard current != -1 else { return nil } + return ComputeUnits(rawValue: current) + } + } + + public func applicationSupportURL() -> URL { + let fileManager = FileManager.default + guard let appDirectoryURL = fileManager.urls(for: .applicationSupportDirectory, in: .userDomainMask).first else { + // To ensure we don't return an optional - if the user domain application support cannot be accessed use the top level application support directory + return URL.applicationSupportDirectory + } + + do { + // Create the application support directory if it doesn't exist + try fileManager.createDirectory(at: appDirectoryURL, withIntermediateDirectories: true, attributes: nil) + return appDirectoryURL + } catch { + print("Error creating application support directory: \(error)") + return fileManager.urls(for: .applicationSupportDirectory, in: .userDomainMask).first! + } + } + + func tempStorageURL() -> URL { + + let tmpDir = applicationSupportURL().appendingPathComponent("hf-diffusion-tmp") + + // Create directory if it doesn't exist + if !FileManager.default.fileExists(atPath: tmpDir.path) { + do { + try FileManager.default.createDirectory(at: tmpDir, withIntermediateDirectories: true, attributes: nil) + } catch { + print("Failed to create temporary directory: \(error)") + return FileManager.default.temporaryDirectory + } + } + + return tmpDir + } + +} diff --git a/Diffusion/Common/Utils.swift b/Diffusion/Common/Utils.swift new file mode 100644 index 0000000..de8dc45 --- /dev/null +++ b/Diffusion/Common/Utils.swift @@ -0,0 +1,101 @@ +// +// Utils.swift +// Diffusion +// +// Created by Pedro Cuenca on 14/1/23. +// See LICENSE at https://github.com/huggingface/swift-coreml-diffusers/LICENSE +// + +import Foundation + +extension String: Error {} + +extension Double { + func formatted(_ format: String) -> String { + return String(format: "\(format)", self) + } +} + +extension String { + var first200Safe: String { + let endIndex = index(startIndex, offsetBy: Swift.min(200, count)) + let substring = String(self[startIndex.. [Bool] { + // Ensure valid parameters + guard numInferenceSteps > 0, numPreviews > 0 else { + return [Bool](repeating: false, count: numInferenceSteps) + } + + // Compute the ideal (floating-point) step size, which represents the average number of steps between previews + let idealStep = Double(numInferenceSteps) / Double(numPreviews) + + // Compute the actual steps at which previews should be made. For each preview, we multiply the ideal step size by the preview number, and round to the nearest integer. + // The result is converted to a `Set` for fast membership tests. + let previewIndices: Set = Set((0.. Double { + let multiplier = pow(10, Double(places)) + let newDecimal = multiplier * self // move the decimal right + let truncated = Double(Int(newDecimal)) // drop the fraction + let originalDecimal = truncated / multiplier // move the decimal back + return originalDecimal + } +} + +func formatLargeNumber(_ n: UInt32) -> String { + let num = abs(Double(n)) + + switch num { + case 1_000_000_000...: + var formatted = num / 1_000_000_000 + formatted = formatted.reduceScale(to: 3) + return "\(formatted)B" + + case 1_000_000...: + var formatted = num / 1_000_000 + formatted = formatted.reduceScale(to: 3) + return "\(formatted)M" + + case 1_000...: + return "\(n)" + + case 0...: + return "\(n)" + + default: + return "\(n)" + } +} diff --git a/Diffusion/Common/Views/PromptTextField.swift b/Diffusion/Common/Views/PromptTextField.swift new file mode 100644 index 0000000..9d20e07 --- /dev/null +++ b/Diffusion/Common/Views/PromptTextField.swift @@ -0,0 +1,141 @@ +// +// PromptTextField.swift +// Diffusion-macOS +// +// Created by Dolmere on 22/06/2023. +// See LICENSE at https://github.com/huggingface/swift-coreml-diffusers/LICENSE +// + +import SwiftUI +import Combine +import StableDiffusion + +struct PromptTextField: View { + @State private var output: String = "" + @State private var input: String = "" + @State private var typing = false + @State private var tokenCount: Int = 0 + @State var isPositivePrompt: Bool = true + @State private var tokenizer: BPETokenizer? + @State private var currentModelVersion: String = "" + + @Binding var textBinding: String + @Binding var model: String // the model version as it's stored in Settings + + private let maxTokenCount = 77 + + private var modelInfo: ModelInfo? { + ModelInfo.from(modelVersion: $model.wrappedValue) + } + + private var pipelineLoader: PipelineLoader? { + guard let modelInfo = modelInfo else { return nil } + return PipelineLoader(model: modelInfo) + } + + private var compiledURL: URL? { + return pipelineLoader?.compiledURL + } + + private var textColor: Color { + switch tokenCount { + case 0...65: + return .green + case 66...75: + return .orange + default: + return .red + } + } + + // macOS initializer + init(text: Binding, isPositivePrompt: Bool, model: Binding) { + _textBinding = text + self.isPositivePrompt = isPositivePrompt + _model = model + } + + // iOS initializer + init(text: Binding, isPositivePrompt: Bool, model: String) { + _textBinding = text + self.isPositivePrompt = isPositivePrompt + _model = .constant(model) + } + + var body: some View { + VStack { + #if os(macOS) + TextField(isPositivePrompt ? "Positive prompt" : "Negative Prompt", text: $textBinding, + axis: .vertical) + .lineLimit(20) + .textFieldStyle(.squareBorder) + .listRowInsets(EdgeInsets(top: 0, leading: -20, bottom: 0, trailing: 20)) + .foregroundColor(textColor == .green ? .primary : textColor) + .frame(minHeight: 30) + if modelInfo != nil && tokenizer != nil { + HStack { + Spacer() + if !textBinding.isEmpty { + Text("\(tokenCount)") + .foregroundColor(textColor) + Text(" / \(maxTokenCount)") + } + } + .onReceive(Just(textBinding)) { text in + updateTokenCount(newText: text) + } + .font(.caption) + } + #else + TextField("Prompt", text: $textBinding, axis: .vertical) + .lineLimit(20) + .listRowInsets(EdgeInsets(top: 0, leading: -20, bottom: 0, trailing: 20)) + .foregroundColor(textColor == .green ? .primary : textColor) + .frame(minHeight: 30) + HStack { + if !textBinding.isEmpty { + Text("\(tokenCount)") + .foregroundColor(textColor) + Text(" / \(maxTokenCount)") + } + Spacer() + } + .onReceive(Just(textBinding)) { text in + updateTokenCount(newText: text) + } + .font(.caption) + #endif + } + .onChange(of: model) { model in + updateTokenCount(newText: textBinding) + } + .onAppear { + updateTokenCount(newText: textBinding) + } + } + + private func updateTokenCount(newText: String) { + // ensure that the compiled URL exists + guard let compiledURL = compiledURL else { return } + // Initialize the tokenizer only when it's not created yet or the model changes + // Check if the model version has changed + let modelVersion = $model.wrappedValue + if modelVersion != currentModelVersion { + do { + tokenizer = try BPETokenizer( + mergesAt: compiledURL.appendingPathComponent("merges.txt"), + vocabularyAt: compiledURL.appendingPathComponent("vocab.json") + ) + currentModelVersion = modelVersion + } catch { + print("Failed to create tokenizer: \(error)") + return + } + } + let (tokens, _) = tokenizer?.tokenize(input: newText) ?? ([], []) + + DispatchQueue.main.async { + self.tokenCount = tokens.count + } + } +} diff --git a/Diffusion/DiffusionApp.swift b/Diffusion/DiffusionApp.swift index d8cdc30..da34543 100644 --- a/Diffusion/DiffusionApp.swift +++ b/Diffusion/DiffusionApp.swift @@ -18,3 +18,13 @@ struct DiffusionApp: App { } let runningOnMac = ProcessInfo.processInfo.isMacCatalystApp +let deviceHas6GBOrMore = ProcessInfo.processInfo.physicalMemory > 5910000000 // Reported by iOS 17 beta (21A5319a) on iPhone 13 Pro: 5917753344 +let deviceHas8GBOrMore = ProcessInfo.processInfo.physicalMemory > 7900000000 // Reported by iOS 17.0.2 on iPhone 15 Pro Max: 8021032960 + +let deviceSupportsQuantization = { + if #available(iOS 17, *) { + true + } else { + false + } +}() diff --git a/Diffusion/DiffusionImage+iOS.swift b/Diffusion/DiffusionImage+iOS.swift new file mode 100644 index 0000000..1abb052 --- /dev/null +++ b/Diffusion/DiffusionImage+iOS.swift @@ -0,0 +1,63 @@ +// +// DiffusionImage+iOS.swift +// Diffusion +// +// Created by Dolmere and Pedro Cuenca on 30/07/2023. +// + +import UIKit +import SwiftUI +import UniformTypeIdentifiers + +extension DiffusionImage { + + /// Instance func to place the generated image on the file system and return the `fileURL` where it is stored. + func save(cgImage: CGImage, filename: String?) -> URL? { + + let image = UIImage(cgImage: cgImage) + let fn = filename ?? "diffusion_generated_image" + let appSupportURL = Settings.shared.tempStorageURL() + + let fileURL = appSupportURL + .appendingPathComponent(fn) + .appendingPathExtension("png") + + if let imageData = image.pngData() { + do { + try imageData.write(to: fileURL) + return fileURL + } catch { + print("Error saving image to temporary file: \(error)") + } + } + return nil + } + + /// Returns a `Data` representation of this generated image in PNG format or nil if there is an error converting the image data. + func pngRepresentation() -> Data? { + let bitmapRep = UIImage(cgImage: cgImage).pngData() + return bitmapRep + } +} + +extension DiffusionImage { + + // MARK: - UIPasteboardWriting + + func writableTypeIdentifiers(for pasteboard: UIPasteboard) -> [String] { + return [UTType.png.identifier] + } + + func itemProviders(forActivityType activityType: UIActivity.ActivityType?) -> [NSItemProvider] { + let itemProvider = NSItemProvider() + itemProvider.registerDataRepresentation(forTypeIdentifier: UTType.png.identifier, visibility: .all) { completion in + guard let pngData = self.pngRepresentation() else { + completion(nil, NSError(domain: "DiffusionImageErrorDomain", code: 0, userInfo: nil)) + return nil + } + completion(pngData, nil) + return nil + } + return [itemProvider] + } +} diff --git a/Diffusion/ModelInfo.swift b/Diffusion/ModelInfo.swift deleted file mode 100644 index 843a853..0000000 --- a/Diffusion/ModelInfo.swift +++ /dev/null @@ -1,129 +0,0 @@ -// -// ModelInfo.swift -// Diffusion -// -// Created by Pedro Cuenca on 29/12/22. -// See LICENSE at https://github.com/huggingface/swift-coreml-diffusers/LICENSE -// - -import CoreML - -enum AttentionVariant: String { - case original - case splitEinsum -} - -extension AttentionVariant { - var defaultComputeUnits: MLComputeUnits { self == .original ? .cpuAndGPU : .cpuAndNeuralEngine } -} - -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 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", splitAttentionSuffix: String = "split_einsum_compiled", supportsEncoder: Bool = false) { - self.modelId = modelId - self.modelVersion = modelVersion - self.originalAttentionSuffix = originalAttentionSuffix - self.splitAttentionSuffix = splitAttentionSuffix - self.supportsEncoder = supportsEncoder - } -} - -extension ModelInfo { - //TODO: set compute units instead and derive variant from it - static var defaultAttention: AttentionVariant { - guard runningOnMac else { return .splitEinsum } - #if os(macOS) - guard Capabilities.hasANE else { return .original } - return Capabilities.performanceCores >= 8 ? .original : .splitEinsum - #else - return .splitEinsum - #endif - } - - static var defaultComputeUnits: MLComputeUnits { defaultAttention.defaultComputeUnits } - - var bestAttention: AttentionVariant { ModelInfo.defaultAttention } - var defaultComputeUnits: MLComputeUnits { bestAttention.defaultComputeUnits } - - 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 - 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 variant for the current platform. - /// Currently using `split_einsum` for iOS and simple performance heuristics for macOS. - var bestURL: URL { modelURL(for: bestAttention) } - - var reduceMemory: Bool { - return !runningOnMac - } -} - -extension ModelInfo { - // TODO: repo does not exist yet - static let v14Base = ModelInfo( - modelId: "pcuenq/coreml-stable-diffusion-1-4", - modelVersion: "CompVis/stable-diffusion-v1-4" - ) - - static let v15Base = ModelInfo( - modelId: "pcuenq/coreml-stable-diffusion-v1-5", - modelVersion: "runwayml/stable-diffusion-v1-5" - ) - - static let v2Base = ModelInfo( - modelId: "pcuenq/coreml-stable-diffusion-2-base", - modelVersion: "stabilityai/stable-diffusion-2-base" - ) - - static let v21Base = ModelInfo( - modelId: "pcuenq/coreml-stable-diffusion-2-1-base", - modelVersion: "stabilityai/stable-diffusion-2-1-base", - supportsEncoder: true - ) - - static let ofaSmall = ModelInfo( - modelId: "pcuenq/coreml-small-stable-diffusion-v0", - modelVersion: "OFA-Sys/small-stable-diffusion-v0" - ) - - static let MODELS = [ - ModelInfo.v14Base, - ModelInfo.v15Base, - ModelInfo.v2Base, - ModelInfo.v21Base, - ModelInfo.ofaSmall - ] - - static func from(modelVersion: String) -> ModelInfo? { - ModelInfo.MODELS.first(where: {$0.modelVersion == modelVersion}) - } - - static func from(modelId: String) -> ModelInfo? { - ModelInfo.MODELS.first(where: {$0.modelId == modelId}) - } -} - -extension ModelInfo : Equatable { - static func ==(lhs: ModelInfo, rhs: ModelInfo) -> Bool { lhs.modelId == rhs.modelId } -} diff --git a/Diffusion/Pipeline/Pipeline.swift b/Diffusion/Pipeline/Pipeline.swift deleted file mode 100644 index 90e17bd..0000000 --- a/Diffusion/Pipeline/Pipeline.swift +++ /dev/null @@ -1,83 +0,0 @@ -// -// Pipeline.swift -// Diffusion -// -// Created by Pedro Cuenca on December 2022. -// See LICENSE at https://github.com/huggingface/swift-coreml-diffusers/LICENSE -// - -import Foundation -import CoreML -import Combine - -import StableDiffusion - -typealias StableDiffusionProgress = StableDiffusionPipeline.Progress - -struct GenerationResult { - var image: CGImage? - var lastSeed: UInt32 - var interval: TimeInterval? - var userCanceled: Bool -} - -class Pipeline { - let pipeline: StableDiffusionPipeline - let maxSeed: UInt32 - - var progress: StableDiffusionProgress? = nil { - didSet { - progressPublisher.value = progress - } - } - lazy private(set) var progressPublisher: CurrentValueSubject = CurrentValueSubject(progress) - - private var canceled = false - - init(_ pipeline: StableDiffusionPipeline, maxSeed: UInt32 = UInt32.max) { - self.pipeline = pipeline - self.maxSeed = maxSeed - } - - func generate( - prompt: String, - negativePrompt: String = "", - scheduler: StableDiffusionScheduler, - numInferenceSteps stepCount: Int = 50, - seed: UInt32? = nil, - guidanceScale: Float = 7.5, - disableSafety: Bool = false - ) throws -> GenerationResult { - let beginDate = Date() - canceled = false - print("Generating...") - let theSeed = seed ?? UInt32.random(in: 0...maxSeed) - let images = try pipeline.generateImages( - prompt: prompt, - negativePrompt: negativePrompt, - imageCount: 1, - stepCount: stepCount, - seed: theSeed, - guidanceScale: guidanceScale, - disableSafety: disableSafety, - scheduler: scheduler - ) { progress in - handleProgress(progress) - return !canceled - } - let interval = Date().timeIntervalSince(beginDate) - print("Got images: \(images) in \(interval)") - - // Unwrap the 1 image we asked for, nil means safety checker triggered - let image = images.compactMap({ $0 }).first - return GenerationResult(image: image, lastSeed: theSeed, interval: interval, userCanceled: canceled) - } - - func handleProgress(_ progress: StableDiffusionPipeline.Progress) { - self.progress = progress - } - - func setCancelled() { - canceled = true - } -} diff --git a/Diffusion/State.swift b/Diffusion/State.swift deleted file mode 100644 index c9b2dfd..0000000 --- a/Diffusion/State.swift +++ /dev/null @@ -1,129 +0,0 @@ -// -// State.swift -// Diffusion -// -// Created by Pedro Cuenca on 17/1/23. -// See LICENSE at https://github.com/huggingface/swift-coreml-diffusers/LICENSE -// - -import Combine -import SwiftUI -import StableDiffusion -import CoreML - -let DEFAULT_MODEL = ModelInfo.v2Base -let DEFAULT_PROMPT = "Labrador in the style of Vermeer" - -enum GenerationState { - case startup - case running(StableDiffusionProgress?) - case complete(String, CGImage?, UInt32, TimeInterval?) - case userCanceled - case failed(Error) -} - -typealias ComputeUnits = MLComputeUnits - -class GenerationContext: ObservableObject { - let scheduler = StableDiffusionScheduler.dpmSolverMultistepScheduler - - @Published var pipeline: Pipeline? = nil { - didSet { - if let pipeline = pipeline { - progressSubscriber = pipeline - .progressPublisher - .receive(on: DispatchQueue.main) - .sink { progress in - guard let progress = progress else { return } - self.state = .running(progress) - } - } - } - } - @Published var state: GenerationState = .startup - - @Published var positivePrompt = DEFAULT_PROMPT - @Published var negativePrompt = "" - - // FIXME: Double to support the slider component - @Published var steps = 25.0 - @Published var numImages = 1.0 - @Published var seed = -1.0 - @Published var guidanceScale = 7.5 - @Published var disableSafety = false - - @Published var computeUnits: ComputeUnits = Settings.shared.userSelectedComputeUnits ?? ModelInfo.defaultComputeUnits - - private var progressSubscriber: Cancellable? - - func generate() async throws -> GenerationResult { - guard let pipeline = pipeline else { throw "No pipeline" } - let seed = self.seed >= 0 ? UInt32(self.seed) : nil - return try pipeline.generate( - prompt: positivePrompt, - negativePrompt: negativePrompt, - scheduler: scheduler, - numInferenceSteps: Int(steps), - seed: seed, - guidanceScale: Float(guidanceScale), - disableSafety: disableSafety - ) - } - - func cancelGeneration() { - pipeline?.setCancelled() - } -} - -class Settings { - static let shared = Settings() - - let defaults = UserDefaults.standard - - enum Keys: String { - case model - case safetyCheckerDisclaimer - case computeUnits - } - - private init() { - defaults.register(defaults: [ - Keys.model.rawValue: ModelInfo.v2Base.modelId, - Keys.safetyCheckerDisclaimer.rawValue: false, - Keys.computeUnits.rawValue: -1 // Use default - ]) - } - - var currentModel: ModelInfo { - set { - defaults.set(newValue.modelId, forKey: Keys.model.rawValue) - } - get { - guard let modelId = defaults.string(forKey: Keys.model.rawValue) else { return DEFAULT_MODEL } - return ModelInfo.from(modelId: modelId) ?? DEFAULT_MODEL - } - } - - var safetyCheckerDisclaimerShown: Bool { - set { - defaults.set(newValue, forKey: Keys.safetyCheckerDisclaimer.rawValue) - } - get { - return defaults.bool(forKey: Keys.safetyCheckerDisclaimer.rawValue) - } - } - - /// Returns the option selected by the user, if overridden - /// `nil` means: guess best - var userSelectedComputeUnits: ComputeUnits? { - set { - // Any value other than the supported ones would cause `get` to return `nil` - defaults.set(newValue?.rawValue ?? -1, forKey: Keys.computeUnits.rawValue) - } - get { - let current = defaults.integer(forKey: Keys.computeUnits.rawValue) - guard current != -1 else { return nil } - return ComputeUnits(rawValue: current) - } - } -} diff --git a/Diffusion/Utils.swift b/Diffusion/Utils.swift deleted file mode 100644 index 480c265..0000000 --- a/Diffusion/Utils.swift +++ /dev/null @@ -1,17 +0,0 @@ -// -// Utils.swift -// Diffusion -// -// Created by Pedro Cuenca on 14/1/23. -// See LICENSE at https://github.com/huggingface/swift-coreml-diffusers/LICENSE -// - -import Foundation - -extension String: Error {} - -extension Double { - func formatted(_ format: String) -> String { - return String(format: "\(format)", self) - } -} diff --git a/Diffusion/Utils_iOS.swift b/Diffusion/Utils_iOS.swift new file mode 100644 index 0000000..9eab45e --- /dev/null +++ b/Diffusion/Utils_iOS.swift @@ -0,0 +1,17 @@ +// +// Utils_iOS.swift +// Diffusion +// +// Created by Dolmere on 31/07/2023. +// + +import SwiftUI + +extension CGImage { + static func fromData(_ imageData: Data) -> CGImage? { + if let image = UIImage(data: imageData)?.cgImage { + return image + } + return nil + } +} diff --git a/Diffusion/Views/Loading.swift b/Diffusion/Views/Loading.swift index a298f56..c64e8f0 100644 --- a/Diffusion/Views/Loading.swift +++ b/Diffusion/Views/Loading.swift @@ -9,9 +9,14 @@ import SwiftUI import Combine -let model = ModelInfo.v2Base +func iosModel() -> ModelInfo { + guard deviceSupportsQuantization else { return ModelInfo.v21Base } + if deviceHas6GBOrMore { return ModelInfo.xlmbpChunked } + return ModelInfo.v21Palettized +} struct LoadingView: View { + @StateObject var generation = GenerationContext() @State private var preparationPhase = "Downloading…" @@ -40,7 +45,7 @@ struct LoadingView: View { .environmentObject(generation) .onAppear { Task.init { - let loader = PipelineLoader(model: model) + let loader = PipelineLoader(model: iosModel()) stateSubscriber = loader.statePublisher.sink { state in DispatchQueue.main.async { switch state { diff --git a/Diffusion/Views/TextToImage.swift b/Diffusion/Views/TextToImage.swift index a97562b..68f1e80 100644 --- a/Diffusion/Views/TextToImage.swift +++ b/Diffusion/Views/TextToImage.swift @@ -52,6 +52,7 @@ struct ShareButtons: View { } struct ImageWithPlaceholder: View { + @EnvironmentObject var generation: GenerationContext var state: Binding var body: some View { @@ -62,10 +63,20 @@ struct ImageWithPlaceholder: View { // The first time it takes a little bit before generation starts return AnyView(ProgressView()) } + let step = Int(progress.step) + 1 let fraction = Double(step) / Double(progress.stepCount) let label = "Step \(step) of \(progress.stepCount)" - return AnyView(ProgressView(label, value: fraction, total: 1).padding()) + return AnyView(VStack { + Group { + if let safeImage = generation.previewImage { + Image(safeImage, scale: 1, label: Text("generated")) + .resizable() + .clipShape(RoundedRectangle(cornerRadius: 20)) + } + } + ProgressView(label, value: fraction, total: 1).padding() + }) case .complete(let lastPrompt, let image, _, let interval): guard let theImage = image else { return AnyView(Image(systemName: "exclamationmark.triangle").resizable()) @@ -113,11 +124,7 @@ struct TextToImage: View { var body: some View { VStack { HStack { - TextField("Prompt", text: $generation.positivePrompt) - .textFieldStyle(.roundedBorder) - .onSubmit { - submit() - } + PromptTextField(text: $generation.positivePrompt, isPositivePrompt: true, model: iosModel().modelVersion) Button("Generate") { submit() } @@ -129,5 +136,6 @@ struct TextToImage: View { Spacer() } .padding() + .environmentObject(generation) } } diff --git a/README.md b/README.md index 30d21cf..ba6f91a 100644 --- a/README.md +++ b/README.md @@ -9,6 +9,8 @@ On first launch, the application downloads a zipped archive with a Core ML versi For faster inference, we use a very fast scheduler: [DPM-Solver++](https://github.com/LuChengTHU/dpm-solver), that we ported to Swift from our [diffusers DPMSolverMultistepScheduler implementation](https://github.com/huggingface/diffusers/blob/main/src/diffusers/schedulers/scheduling_dpmsolver_multistep.py). +The app supports models quantized with `coremltools` version 7 or better. This requires macOS 14 or iOS/iPadOS 17. + ## Compatibility and Performance - macOS Ventura 13.1, iOS/iPadOS 16.2, Xcode 14.2. @@ -18,6 +20,8 @@ For faster inference, we use a very fast scheduler: [DPM-Solver++](https://githu See [this post](https://huggingface.co/blog/fast-mac-diffusers) and [this issue](https://github.com/huggingface/swift-coreml-diffusers/issues/31) for additional performance figures. +Quantized models run faster, but they require macOS Ventura 14, or iOS/iPadOS 17. + The application will try to guess the best hardware to run models on. You can override this setting using the `Advanced` section in the controls sidebar. ## How to Run diff --git a/config/common.xcconfig b/config/common.xcconfig index 8c20eb6..f582cb1 100644 --- a/config/common.xcconfig +++ b/config/common.xcconfig @@ -10,8 +10,8 @@ // https://help.apple.com/xcode/#/dev745c5c974 PRODUCT_NAME = Diffusers -CURRENT_PROJECT_VERSION = 1.1.0 -MARKETING_VERSION = 1.1 +CURRENT_PROJECT_VERSION = 20240607.135124 +MARKETING_VERSION = 1.5 // Update if you fork this repo DEVELOPMENT_TEAM = 2EADP68M95