DiffusionImage

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
+    }
+}

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
+    }
+}

Diffusion.xcodeproj/project.pbxproj

@@ -7,8 +7,14 @@
 	objects = {
 
 /* Begin PBXBuildFile section */
+		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 */; };
 		EB25B3D62A3A2DC4000E25A1 /* StableDiffusion in Frameworks */ = {isa = PBXBuildFile; productRef = EB25B3D52A3A2DC4000E25A1 /* StableDiffusion */; };
 		EB25B3D82A3A2DD5000E25A1 /* StableDiffusion in Frameworks */ = {isa = PBXBuildFile; productRef = EB25B3D72A3A2DD5000E25A1 /* StableDiffusion */; };
@@ -63,7 +69,12 @@
 /* End PBXContainerItemProxy section */
 
 /* Begin PBXFileReference section */
+		8C4B32032A770C1D0090EF17 /* DiffusionImage+macOS.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = "DiffusionImage+macOS.swift"; sourceTree = "<group>"; };
+		8C4B32052A770C300090EF17 /* DiffusionImage+iOS.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = "DiffusionImage+iOS.swift"; sourceTree = "<group>"; };
+		8C4B32072A77F90C0090EF17 /* Utils_iOS.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = Utils_iOS.swift; sourceTree = "<group>"; };
+		8C4B32092A77F9160090EF17 /* Utils_macOS.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = Utils_macOS.swift; sourceTree = "<group>"; };
 		8CD8A5392A456EF800BD8A98 /* PromptTextField.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = PromptTextField.swift; sourceTree = "<group>"; };
+		8CEEB7D82A54C88C00C23829 /* DiffusionImage.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = DiffusionImage.swift; sourceTree = "<group>"; };
 		EB067F862992E561004D1AD9 /* HelpContent.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = HelpContent.swift; sourceTree = "<group>"; };
 		EB33A51E2954E1BC00B16357 /* Info.plist */ = {isa = PBXFileReference; lastKnownFileType = text.plist; path = Info.plist; sourceTree = "<group>"; };
 		EB560F0329A3C20800C0F8B8 /* Capabilities.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = Capabilities.swift; sourceTree = "<group>"; };
@@ -152,6 +163,7 @@
 				EBE3FF4B295E1EFE00E921AA /* ModelInfo.swift */,
 				EBDD7DB72976AAFE00C1C4B2 /* State.swift */,
 				EBDD7DB22973200200C1C4B2 /* Utils.swift */,
+				8CEEB7D82A54C88C00C23829 /* DiffusionImage.swift */,
 				EBB5BA5129425B07003A2A5B /* Pipeline */,
 				8CD8A53B2A476E1C00BD8A98 /* Views */,
 			);
@@ -206,6 +218,8 @@
 				EBE755CC293E37DD00806B32 /* Assets.xcassets */,
 				EBE755CE293E37DD00806B32 /* Diffusion.entitlements */,
 				EBE755CF293E37DD00806B32 /* Preview Content */,
+				8C4B32052A770C300090EF17 /* DiffusionImage+iOS.swift */,
+				8C4B32072A77F90C0090EF17 /* Utils_iOS.swift */,
 			);
 			path = Diffusion;
 			sourceTree = "<group>";
@@ -269,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 */,
 			);
@@ -484,8 +500,11 @@
 				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 */,
@@ -520,6 +539,7 @@
 				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 */,
@@ -528,7 +548,9 @@
 				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;
 		};

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
+    }
+}

Diffusion/Common/Pipeline/Pipeline.swift

@@ -80,7 +80,7 @@ class Pipeline {
         config.seed = theSeed
         config.guidanceScale = guidanceScale
         config.disableSafety = disableSafety
-        config.schedulerType = scheduler
+        config.schedulerType = scheduler.asStableDiffusionScheduler()
         config.useDenoisedIntermediates = true
 
         // Evenly distribute previews based on inference steps