macOS UI

.gitignore

@@ -1 +1,2 @@
 .DS_Store
+xcuserdata/

Diffusion-macOS/ContentView.swift

@@ -0,0 +1,100 @@
+//
+//  ContentView.swift
+//  Diffusion-macOS
+//
+//  Created by Cyril Zakka on 1/12/23.
+//  See LICENSE at https://github.com/huggingface/swift-coreml-diffusers/LICENSE
+//
+
+import SwiftUI
+import ImageIO
+
+
+// AppKit version that uses NSImage, NSSavePanel
+struct ShareButtons: View {
+    var image: CGImage
+    var name: String
+
+    var filename: String {
+        name.replacingOccurrences(of: " ", with: "_")
+    }
+
+    func showSavePanel() -> URL? {
+        let savePanel = NSSavePanel()
+        savePanel.allowedContentTypes = [.png]
+        savePanel.canCreateDirectories = true
+        savePanel.isExtensionHidden = false
+        savePanel.title = "Save your image"
+        savePanel.message = "Choose a folder and a name to store the image."
+        savePanel.nameFieldLabel = "File name:"
+        savePanel.nameFieldStringValue = filename
+
+        let response = savePanel.runModal()
+        return response == .OK ? savePanel.url : nil
+    }
+
+    func savePNG(cgImage: CGImage, path: URL) {
+        let image = NSImage(cgImage: cgImage, size: .zero)
+        let imageRepresentation = NSBitmapImageRep(data: image.tiffRepresentation!)
+        guard let pngData = imageRepresentation?.representation(using: .png, properties: [:]) else {
+            print("Error generating PNG data")
+            return
+        }
+        do {
+            try pngData.write(to: path)
+        } catch {
+            print("Error saving: \(error)")
+        }
+    }
+
+    var body: some View {
+        let imageView = Image(image, scale: 1, label: Text(name))
+        HStack {
+            ShareLink(item: imageView, preview: SharePreview(name, image: imageView))
+            Button() {
+                if let url = showSavePanel() {
+                    savePNG(cgImage: image, path: url)
+                }
+            } label: {
+                Label("Save…", systemImage: "square.and.arrow.down")
+            }
+        }
+    }
+}
+
+struct ContentView: View {
+    @StateObject var generation = GenerationContext()
+
+    func toolbar() -> any View {
+        if case .complete(let prompt, let cgImage, _) = generation.state, let cgImage = cgImage {
+            return ShareButtons(image: cgImage, name: prompt)
+        } else {
+            let prompt = DEFAULT_PROMPT
+            let cgImage = NSImage(imageLiteralResourceName: "placeholder").cgImage(forProposedRect: nil, context: nil, hints: nil)!
+            return ShareButtons(image: cgImage, name: prompt)
+        }
+    }
+
+    var body: some View {
+        NavigationSplitView {
+            ControlsView()
+                .navigationSplitViewColumnWidth(min: 250, ideal: 300)
+        } detail: {
+            GeneratedImageView()
+                .aspectRatio(contentMode: .fit)
+                .frame(width: 512, height: 512)
+                .cornerRadius(15)
+                .toolbar {
+                    AnyView(toolbar())
+                }
+
+        }
+        .environmentObject(generation)
+    }
+}
+
+struct ContentView_Previews: PreviewProvider {
+    static var previews: some View {
+        ContentView()
+    }
+}

Diffusion-macOS/ControlsView.swift

@@ -0,0 +1,151 @@
+//
+//  PromptView.swift
+//  Diffusion-macOS
+//
+//  Created by Cyril Zakka on 1/12/23.
+//  See LICENSE at https://github.com/huggingface/swift-coreml-diffusers/LICENSE
+//
+
+import Combine
+import SwiftUI
+import CompactSlider
+
+enum PipelineState {
+    case downloading(Double)
+    case uncompressing
+    case loading
+    case ready
+    case failed(Error)
+}
+
+struct ControlsView: View {
+    @EnvironmentObject var generation: GenerationContext
+
+    static let models = ModelInfo.MODELS
+    static let modelNames = models.map { $0.modelVersion }
+
+    @State private var model = Settings.shared.currentModel.modelVersion
+    @State private var disclosedPrompt = true
+
+    // TODO: refactor download with similar code in Loading.swift (iOS)
+    @State private var stateSubscriber: Cancellable?
+    @State private var pipelineState: PipelineState = .downloading(0)
+
+    func modelDidChange(model: ModelInfo) {
+        print("Loading model \(model)")
+        Settings.shared.currentModel = model
+
+        pipelineState = .downloading(0)
+        Task.init {
+            let loader = PipelineLoader(model: model)
+            stateSubscriber = loader.statePublisher.sink { state in
+                DispatchQueue.main.async {
+                    switch state {
+                    case .downloading(let progress):
+                        pipelineState = .downloading(progress)
+                    case .uncompressing:
+                        pipelineState = .uncompressing
+                    case .readyOnDisk:
+                        pipelineState = .loading
+                    default:
+                        break
+                    }
+                }
+            }
+            do {
+                generation.pipeline = try await loader.prepare()
+                pipelineState = .ready
+            } catch {
+                print("Could not load model, error: \(error)")
+                pipelineState = .failed(error)
+            }
+        }
+    }
+
+    var body: some View {
+        VStack(alignment: .leading) {
+
+            Label("Adjustments", systemImage: "gearshape.2")
+                .font(.headline)
+                .fontWeight(.bold)
+            Divider()
+
+            ScrollView {
+                Group {
+                    DisclosureGroup {
+                        Picker("", selection: $model) {
+                            ForEach(Self.modelNames, id: \.self) {
+                                Text($0)
+                            }
+                        }
+                        .onChange(of: model) { theModel in
+                            guard let model = ModelInfo.from(modelVersion: theModel) else { return }
+                            modelDidChange(model: model)
+                        }
+                    } label: {
+                        Label("Model", systemImage: "cpu").foregroundColor(.secondary)
+                    }
+
+                    Divider()
+
+                    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)
+                        }.padding(.leading, 10)
+                    } label: {
+                        Label("Prompts", systemImage: "text.quote").foregroundColor(.secondary)
+                    }
+
+                    Divider()
+
+                    DisclosureGroup {
+                        CompactSlider(value: $generation.steps, in: 0...150, step: 5) {
+                            Text("Steps")
+                            Spacer()
+                            Text("\(Int(generation.steps))")
+                        }.padding(.leading, 10)
+                    } label: {
+                        Label("Step count", systemImage: "square.3.layers.3d.down.left").foregroundColor(.secondary)
+                    }
+                    Divider()
+
+//                    DisclosureGroup() {
+//                        CompactSlider(value: $generation.numImages, in: 0...10, step: 1) {
+//                            Text("Number of Images")
+//                            Spacer()
+//                            Text("\(Int(generation.numImages))")
+//                        }.padding(.leading, 10)
+//                    } label: {
+//                        Label("Number of images", systemImage: "photo.stack").foregroundColor(.secondary)
+//                    }
+//                    Divider()
+
+                    DisclosureGroup() {
+                        let sliderLabel = generation.seed < 0 ? "Random Seed" : "Seed"
+                        CompactSlider(value: $generation.seed, in: -1...1000, step: 1) {
+                            Text(sliderLabel)
+                            Spacer()
+                            Text("\(Int(generation.seed))")
+                        }.padding(.leading, 10)
+                    } label: {
+                        Label("Seed", systemImage: "leaf").foregroundColor(.secondary)
+                    }
+                }
+            }
+
+            StatusView(pipelineState: $pipelineState)
+        }
+        .padding()
+        .onAppear {
+            print(PipelineLoader.models)
+            modelDidChange(model: ModelInfo.from(modelVersion: model) ?? ModelInfo.v2Base)
+        }
+    }
+}
+

Diffusion-macOS/Diffusion_macOS.entitlements

@@ -0,0 +1,12 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<!DOCTYPE plist PUBLIC "-//Apple//DTD PLIST 1.0//EN" "http://www.apple.com/DTDs/PropertyList-1.0.dtd">
+<plist version="1.0">
+<dict>
+	<key>com.apple.security.app-sandbox</key>
+	<true/>
+	<key>com.apple.security.files.user-selected.read-write</key>
+	<true/>
+	<key>com.apple.security.network.client</key>
+	<true/>
+</dict>
+</plist>

Diffusion-macOS/Diffusion_macOSApp.swift

@@ -0,0 +1,20 @@
+//
+//  Diffusion_macOSApp.swift
+//  Diffusion-macOS
+//
+//  Created by Cyril Zakka on 1/12/23.
+//  See LICENSE at https://github.com/huggingface/swift-coreml-diffusers/LICENSE
+//
+
+import SwiftUI
+
+@main
+struct Diffusion_macOSApp: App {
+    var body: some Scene {
+        WindowGroup {
+            ContentView()
+        }
+    }
+}
+
+let runningOnMac = true

Diffusion-macOS/GeneratedImageView.swift

@@ -0,0 +1,37 @@
+//
+//  GeneratedImageView.swift
+//  Diffusion
+//
+//  Created by Pedro Cuenca on 18/1/23.
+//  See LICENSE at https://github.com/huggingface/swift-coreml-diffusers/LICENSE
+//
+
+import SwiftUI
+
+struct GeneratedImageView: View {
+    @EnvironmentObject var generation: GenerationContext
+
+    var body: some View {
+        switch generation.state {
+        case .startup: return AnyView(Image("placeholder").resizable())
+        case .running(let progress):
+            guard let progress = progress, progress.stepCount > 0 else {
+                // 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())
+        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))
+            )
+        }
+    }
+}

Diffusion-macOS/Preview Content/Preview Assets.xcassets/Contents.json

@@ -0,0 +1,6 @@
+{
+  "info" : {
+    "author" : "xcode",
+    "version" : 1
+  }
+}