Skip to content

Commit 4eab476

Browse files
dolmerepcuenca
dolmere
and
pcuenca
authored
DiffusionImage (huggingface#64)
* DiffusionImage Addition of DiffusionImage Model extracted from 59 huggingface#59 * Update Latest DiffusionImage updates to support drag and drop and copy. * Update Diffusion/Common/DiffusionImage.swift Docc compatible comment style. Co-authored-by: Pedro Cuenca <pedro@huggingface.co> * Update Diffusion/Common/DiffusionImage.swift typo fixes in descriptive note Co-authored-by: Pedro Cuenca <pedro@huggingface.co> * Better comments Fleshing out docc compatible comments in new class. Also refactored two identical functions with different names. * Manual merge fix Fix from recent manual merge * Custom local enum created Diffusion_StableDiffusionScheduler created to get around a problem in a future PR for copy/paste and drag/drop with preserved "recipe" data about how the image was created. * Review suggestions This is a consolidated commit with the suggestions submitted in review. - Move as much platform specific code as possible out into target specific files. - Move save() and pngRepresentation() into separate platform specific files - Move platform specific protocol code out into separate file. * Less platform checks New Utils_iOS and Utils_macOS allow platform specific utility functions to live inside the target file structure. - added CGImage.fromData(Data) funcs for macOS/iOS - added new func inside StableDiffusionScheduler enum. - renamed Diffusion_StableDiffusionScheduler to StableDiffusionScheduler * Update Diffusion/Common/DiffusionImage.swift remove comment with answered question Co-authored-by: Pedro Cuenca <pedro@huggingface.co> * Update Diffusion-macOS/DiffusionImage+macOS.swift remove comment from self descriptive code Co-authored-by: Pedro Cuenca <pedro@huggingface.co> * Update Diffusion-macOS/DiffusionImage+macOS.swift remove comment from self descriptive code Co-authored-by: Pedro Cuenca <pedro@huggingface.co> * Update Diffusion/Common/DiffusionImage.swift more concise comment Co-authored-by: Pedro Cuenca <pedro@huggingface.co> --------- Co-authored-by: Pedro Cuenca <pedro@huggingface.co>
1 parent e5f9a65 commit 4eab476

9 files changed

+417
-2
lines changed

‎Diffusion-macOS/DiffusionImage+macOS.swift

+112
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,112 @@
1+
//
2+
// DiffusionImage+macOS.swift
3+
// Diffusion-macOS
4+
//
5+
// Created by Dolmere and Pedro Cuenca on 30/07/2023.
6+
//
7+
8+
import SwiftUI
9+
import UniformTypeIdentifiers
10+
11+
extension DiffusionImage {
12+
13+
/// Instance func to place the generated image on the file system and return the `fileURL` where it is stored.
14+
func save(cgImage: CGImage, filename: String?) -> URL? {
15+
16+
let nsImage = NSImage(cgImage: cgImage, size: NSSize(width: cgImage.width, height: cgImage.height))
17+
18+
19+
let appSupportURL = Settings.shared.tempStorageURL()
20+
let fn = filename ?? "diffusion_generated_image"
21+
let fileURL = appSupportURL
22+
.appendingPathComponent(fn)
23+
.appendingPathExtension("png")
24+
25+
// Save the image as a temporary file
26+
if let tiffData = nsImage.tiffRepresentation,
27+
let bitmap = NSBitmapImageRep(data: tiffData),
28+
let pngData = bitmap.representation(using: .png, properties: [:]) {
29+
do {
30+
try pngData.write(to: fileURL)
31+
return fileURL
32+
} catch {
33+
print("Error saving image to temporary file: \(error)")
34+
}
35+
}
36+
return nil
37+
}
38+
39+
/// Returns a `Data` representation of this generated image in PNG format or nil if there is an error converting the image data.
40+
func pngRepresentation() -> Data? {
41+
let bitmapRep = NSBitmapImageRep(cgImage: cgImage)
42+
return bitmapRep.representation(using: .png, properties: [:])
43+
}
44+
}
45+
46+
extension DiffusionImage: NSItemProviderWriting {
47+
48+
// MARK: - NSItemProviderWriting
49+
50+
static var writableTypeIdentifiersForItemProvider: [String] {
51+
return [UTType.data.identifier, UTType.png.identifier, UTType.fileURL.identifier]
52+
}
53+
54+
func itemProviderVisibilityForRepresentation(withTypeIdentifier typeIdentifier: String) -> NSItemProviderRepresentationVisibility {
55+
return .all
56+
}
57+
58+
func itemProviderRepresentation(forTypeIdentifier typeIdentifier: String) throws -> NSItemProvider {
59+
print("itemProviderRepresentation(forTypeIdentifier")
60+
print(typeIdentifier)
61+
let data = try NSKeyedArchiver.archivedData(withRootObject: self, requiringSecureCoding: true)
62+
let itemProvider = NSItemProvider()
63+
itemProvider.registerDataRepresentation(forTypeIdentifier: typeIdentifier, visibility: NSItemProviderRepresentationVisibility.all) { completion in
64+
completion(data, nil)
65+
return nil
66+
}
67+
return itemProvider
68+
}
69+
70+
func loadData(withTypeIdentifier typeIdentifier: String, forItemProviderCompletionHandler completionHandler: @escaping @Sendable (Data?, Error?) -> Void) -> Progress? {
71+
if typeIdentifier == NSPasteboard.PasteboardType.fileURL.rawValue {
72+
let data = fileURL.dataRepresentation
73+
completionHandler(data, nil)
74+
} else if typeIdentifier == UTType.png.identifier {
75+
let data = pngRepresentation()
76+
completionHandler(data, nil)
77+
} else {
78+
// Indicate that the specified typeIdentifier is not supported
79+
let error = NSError(domain: "com.huggingface.diffusion", code: 0, userInfo: [NSLocalizedDescriptionKey: "Unsupported typeIdentifier"])
80+
completionHandler(nil, error)
81+
}
82+
return nil
83+
}
84+
85+
}
86+
87+
extension DiffusionImage: NSPasteboardWriting {
88+
89+
// MARK: - NSPasteboardWriting
90+
91+
func writableTypes(for pasteboard: NSPasteboard) -> [NSPasteboard.PasteboardType] {
92+
return [
93+
NSPasteboard.PasteboardType.fileURL,
94+
NSPasteboard.PasteboardType(rawValue: UTType.png.identifier)
95+
]
96+
}
97+
98+
func pasteboardPropertyList(forType type: NSPasteboard.PasteboardType) -> Any? {
99+
if type == NSPasteboard.PasteboardType.fileURL {
100+
101+
// Return the file's data' representation
102+
return fileURL.dataRepresentation
103+
104+
} else if type.rawValue == UTType.png.identifier {
105+
106+
// Return a PNG data representation
107+
return pngRepresentation()
108+
}
109+
110+
return nil
111+
}
112+
}

‎Diffusion-macOS/Utils_macOS.swift

+17
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,17 @@
1+
//
2+
// Utils_macOS.swift
3+
// Diffusion-macOS
4+
//
5+
// Created by Dolmere on 31/07/2023.
6+
//
7+
8+
import SwiftUI
9+
10+
extension CGImage {
11+
static func fromData(_ imageData: Data) -> CGImage? {
12+
if let image = NSBitmapImageRep(data: imageData)?.cgImage {
13+
return image
14+
}
15+
return nil
16+
}
17+
}

‎Diffusion.xcodeproj/project.pbxproj

+22
Original file line numberDiff line numberDiff line change
@@ -7,8 +7,14 @@
77
objects = {
88

99
/* Begin PBXBuildFile section */
10+
8C4B32042A770C1D0090EF17 /* DiffusionImage+macOS.swift in Sources */ = {isa = PBXBuildFile; fileRef = 8C4B32032A770C1D0090EF17 /* DiffusionImage+macOS.swift */; };
11+
8C4B32062A770C300090EF17 /* DiffusionImage+iOS.swift in Sources */ = {isa = PBXBuildFile; fileRef = 8C4B32052A770C300090EF17 /* DiffusionImage+iOS.swift */; };
12+
8C4B32082A77F90C0090EF17 /* Utils_iOS.swift in Sources */ = {isa = PBXBuildFile; fileRef = 8C4B32072A77F90C0090EF17 /* Utils_iOS.swift */; };
13+
8C4B320A2A77F9160090EF17 /* Utils_macOS.swift in Sources */ = {isa = PBXBuildFile; fileRef = 8C4B32092A77F9160090EF17 /* Utils_macOS.swift */; };
1014
8CD8A53A2A456EF800BD8A98 /* PromptTextField.swift in Sources */ = {isa = PBXBuildFile; fileRef = 8CD8A5392A456EF800BD8A98 /* PromptTextField.swift */; };
1115
8CD8A53C2A476E2C00BD8A98 /* PromptTextField.swift in Sources */ = {isa = PBXBuildFile; fileRef = 8CD8A5392A456EF800BD8A98 /* PromptTextField.swift */; };
16+
8CEEB7D92A54C88C00C23829 /* DiffusionImage.swift in Sources */ = {isa = PBXBuildFile; fileRef = 8CEEB7D82A54C88C00C23829 /* DiffusionImage.swift */; };
17+
8CEEB7DA2A54C88C00C23829 /* DiffusionImage.swift in Sources */ = {isa = PBXBuildFile; fileRef = 8CEEB7D82A54C88C00C23829 /* DiffusionImage.swift */; };
1218
EB067F872992E561004D1AD9 /* HelpContent.swift in Sources */ = {isa = PBXBuildFile; fileRef = EB067F862992E561004D1AD9 /* HelpContent.swift */; };
1319
EB25B3D62A3A2DC4000E25A1 /* StableDiffusion in Frameworks */ = {isa = PBXBuildFile; productRef = EB25B3D52A3A2DC4000E25A1 /* StableDiffusion */; };
1420
EB25B3D82A3A2DD5000E25A1 /* StableDiffusion in Frameworks */ = {isa = PBXBuildFile; productRef = EB25B3D72A3A2DD5000E25A1 /* StableDiffusion */; };
@@ -63,7 +69,12 @@
6369
/* End PBXContainerItemProxy section */
6470

6571
/* Begin PBXFileReference section */
72+
8C4B32032A770C1D0090EF17 /* DiffusionImage+macOS.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = "DiffusionImage+macOS.swift"; sourceTree = "<group>"; };
73+
8C4B32052A770C300090EF17 /* DiffusionImage+iOS.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = "DiffusionImage+iOS.swift"; sourceTree = "<group>"; };
74+
8C4B32072A77F90C0090EF17 /* Utils_iOS.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = Utils_iOS.swift; sourceTree = "<group>"; };
75+
8C4B32092A77F9160090EF17 /* Utils_macOS.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = Utils_macOS.swift; sourceTree = "<group>"; };
6676
8CD8A5392A456EF800BD8A98 /* PromptTextField.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = PromptTextField.swift; sourceTree = "<group>"; };
77+
8CEEB7D82A54C88C00C23829 /* DiffusionImage.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = DiffusionImage.swift; sourceTree = "<group>"; };
6778
EB067F862992E561004D1AD9 /* HelpContent.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = HelpContent.swift; sourceTree = "<group>"; };
6879
EB33A51E2954E1BC00B16357 /* Info.plist */ = {isa = PBXFileReference; lastKnownFileType = text.plist; path = Info.plist; sourceTree = "<group>"; };
6980
EB560F0329A3C20800C0F8B8 /* Capabilities.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = Capabilities.swift; sourceTree = "<group>"; };
@@ -152,6 +163,7 @@
152163
EBE3FF4B295E1EFE00E921AA /* ModelInfo.swift */,
153164
EBDD7DB72976AAFE00C1C4B2 /* State.swift */,
154165
EBDD7DB22973200200C1C4B2 /* Utils.swift */,
166+
8CEEB7D82A54C88C00C23829 /* DiffusionImage.swift */,
155167
EBB5BA5129425B07003A2A5B /* Pipeline */,
156168
8CD8A53B2A476E1C00BD8A98 /* Views */,
157169
);
@@ -206,6 +218,8 @@
206218
EBE755CC293E37DD00806B32 /* Assets.xcassets */,
207219
EBE755CE293E37DD00806B32 /* Diffusion.entitlements */,
208220
EBE755CF293E37DD00806B32 /* Preview Content */,
221+
8C4B32052A770C300090EF17 /* DiffusionImage+iOS.swift */,
222+
8C4B32072A77F90C0090EF17 /* Utils_iOS.swift */,
209223
);
210224
path = Diffusion;
211225
sourceTree = "<group>";
@@ -269,6 +283,8 @@
269283
F155203329710B3600DC009B /* StatusView.swift */,
270284
EB067F862992E561004D1AD9 /* HelpContent.swift */,
271285
EB560F0329A3C20800C0F8B8 /* Capabilities.swift */,
286+
8C4B32032A770C1D0090EF17 /* DiffusionImage+macOS.swift */,
287+
8C4B32092A77F9160090EF17 /* Utils_macOS.swift */,
272288
F155202C2971093400DC009B /* Diffusion_macOS.entitlements */,
273289
F15520292971093400DC009B /* Preview Content */,
274290
);
@@ -484,8 +500,11 @@
484500
EBE75602293E91E200806B32 /* Pipeline.swift in Sources */,
485501
EBE755CB293E37DD00806B32 /* TextToImage.swift in Sources */,
486502
EBB5BA5A29426E06003A2A5B /* Downloader.swift in Sources */,
503+
8C4B32062A770C300090EF17 /* DiffusionImage+iOS.swift in Sources */,
504+
8CEEB7D92A54C88C00C23829 /* DiffusionImage.swift in Sources */,
487505
EBE3FF4C295E1EFE00E921AA /* ModelInfo.swift in Sources */,
488506
EBE756092941178600806B32 /* Loading.swift in Sources */,
507+
8C4B32082A77F90C0090EF17 /* Utils_iOS.swift in Sources */,
489508
EBDD7DB82976AAFE00C1C4B2 /* State.swift in Sources */,
490509
EBB5BA5329425BEE003A2A5B /* PipelineLoader.swift in Sources */,
491510
8CD8A53C2A476E2C00BD8A98 /* PromptTextField.swift in Sources */,
@@ -520,6 +539,7 @@
520539
F15520262971093300DC009B /* ContentView.swift in Sources */,
521540
EBDD7DB92976AAFE00C1C4B2 /* State.swift in Sources */,
522541
EB067F872992E561004D1AD9 /* HelpContent.swift in Sources */,
542+
8C4B320A2A77F9160090EF17 /* Utils_macOS.swift in Sources */,
523543
EBDD7DB42973200200C1C4B2 /* Utils.swift in Sources */,
524544
8CD8A53A2A456EF800BD8A98 /* PromptTextField.swift in Sources */,
525545
F1552031297109C300DC009B /* ControlsView.swift in Sources */,
@@ -528,7 +548,9 @@
528548
EB560F0429A3C20800C0F8B8 /* Capabilities.swift in Sources */,
529549
F15520242971093300DC009B /* Diffusion_macOSApp.swift in Sources */,
530550
EBDD7DB52973201800C1C4B2 /* ModelInfo.swift in Sources */,
551+
8C4B32042A770C1D0090EF17 /* DiffusionImage+macOS.swift in Sources */,
531552
EBDD7DBD2977FFB300C1C4B2 /* GeneratedImageView.swift in Sources */,
553+
8CEEB7DA2A54C88C00C23829 /* DiffusionImage.swift in Sources */,
532554
);
533555
runOnlyForDeploymentPostprocessing = 0;
534556
};

‎Diffusion/Common/DiffusionImage.swift

+131
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,131 @@
1+
//
2+
// DiffusionImage.swift
3+
// Diffusion
4+
//
5+
// Created by Dolmere on 03/07/2023.
6+
//
7+
8+
import SwiftUI
9+
import StableDiffusion
10+
import CoreTransferable
11+
12+
/// Tracking for a `DiffusionImage` generation state.
13+
enum DiffusionImageState {
14+
case generating
15+
case waiting
16+
case complete
17+
}
18+
19+
/// Generic custom error to use when an image generation fails.
20+
enum DiffusionImageError: Error {
21+
case invalidDiffusionImage
22+
}
23+
24+
/// Combination of a `DiffusionImage` and its associated `DiffusionImageState`
25+
struct DiffusionImageWrapper {
26+
var diffusionImageState: DiffusionImageState = .waiting
27+
var diffusionImage: DiffusionImage? = nil
28+
}
29+
30+
/// Model class to hold a generated image and the "recipe" data that was used to generate it
31+
final class DiffusionImage: NSObject, Identifiable, NSCoding, NSSecureCoding {
32+
33+
let id: UUID
34+
let cgImage: CGImage
35+
let seed: UInt32
36+
let steps: Double
37+
let positivePrompt: String
38+
let negativePrompt: String
39+
let guidanceScale: Double
40+
let disableSafety: Bool
41+
/// Local enum represented with a String to conform to NSSecureCoding
42+
let scheduler: StableDiffusionScheduler
43+
44+
/// 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.
45+
var generatedFilename: String {
46+
return "\(seed)-\(positivePrompt)".first200Safe
47+
}
48+
49+
/// The location on the file system where this generated image is stored.
50+
var fileURL: URL
51+
52+
init(id: UUID, cgImage: CGImage, seed: UInt32, steps: Double, positivePrompt: String, negativePrompt: String, guidanceScale: Double, disableSafety: Bool, scheduler: StableDiffusionScheduler) {
53+
let genname = "\(seed)-\(positivePrompt)".first200Safe
54+
self.id = id
55+
self.cgImage = cgImage
56+
self.seed = seed
57+
self.steps = steps
58+
self.positivePrompt = positivePrompt
59+
self.negativePrompt = negativePrompt
60+
self.guidanceScale = guidanceScale
61+
self.disableSafety = disableSafety
62+
self.scheduler = scheduler
63+
// 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.
64+
self.fileURL = URL.applicationDirectory
65+
// init the instance fully before executing an instance function
66+
super.init()
67+
if let url = save(cgImage: cgImage, filename: genname) {
68+
self.fileURL = url
69+
} else {
70+
fatalError("Fatal error init of DiffusionImage, cannot create image file at \(genname)")
71+
}
72+
}
73+
74+
func encode(with coder: NSCoder) {
75+
coder.encode(id, forKey: "id")
76+
coder.encode(seed, forKey: "seed")
77+
coder.encode(steps, forKey: "steps")
78+
coder.encode(positivePrompt, forKey: "positivePrompt")
79+
coder.encode(negativePrompt, forKey: "negativePrompt")
80+
coder.encode(guidanceScale, forKey: "guidanceScale")
81+
coder.encode(disableSafety, forKey: "disableSafety")
82+
coder.encode(scheduler, forKey: "scheduler")
83+
// Encode cgImage as data
84+
if let data = pngRepresentation() {
85+
coder.encode(data, forKey: "cgImage")
86+
}
87+
}
88+
89+
required init?(coder: NSCoder) {
90+
guard let id = coder.decodeObject(forKey: "id") as? UUID else {
91+
return nil
92+
}
93+
94+
self.id = id
95+
self.seed = UInt32(coder.decodeInt32(forKey: "seed"))
96+
self.steps = coder.decodeDouble(forKey: "steps")
97+
self.positivePrompt = coder.decodeObject(forKey: "positivePrompt") as? String ?? ""
98+
self.negativePrompt = coder.decodeObject(forKey: "negativePrompt") as? String ?? ""
99+
self.guidanceScale = coder.decodeDouble(forKey: "guidanceScale")
100+
self.disableSafety = coder.decodeBool(forKey: "disableSafety")
101+
self.scheduler = coder.decodeObject(forKey: "scheduler") as? StableDiffusionScheduler ?? StableDiffusionScheduler.dpmSolverMultistepScheduler
102+
let genname = "\(seed)-\(positivePrompt)".first200Safe
103+
104+
// Decode cgImage from data
105+
if let imageData = coder.decodeObject(forKey: "cgImage") as? Data {
106+
guard let img = CGImage.fromData(imageData) else { fatalError("Fatal error loading data with missing cgImage in object") }
107+
self.cgImage = img
108+
} else {
109+
fatalError("Fatal error loading data with missing cgImage in object")
110+
}
111+
self.fileURL = URL.applicationDirectory
112+
super.init()
113+
if let url = save(cgImage: cgImage, filename: genname) {
114+
self.fileURL = url
115+
} else {
116+
fatalError("Fatal error init of DiffusionImage, cannot create image file at \(genname)")
117+
}
118+
}
119+
120+
// MARK: - Equatable
121+
122+
static func == (lhs: DiffusionImage, rhs: DiffusionImage) -> Bool {
123+
return lhs.id == rhs.id
124+
}
125+
126+
// MARK: - NSSecureCoding
127+
128+
static var supportsSecureCoding: Bool {
129+
return true
130+
}
131+
}

‎Diffusion/Common/Pipeline/Pipeline.swift

+1-1
Original file line numberDiff line numberDiff line change
@@ -79,7 +79,7 @@ class Pipeline {
7979
config.seed = theSeed
8080
config.guidanceScale = guidanceScale
8181
config.disableSafety = disableSafety
82-
config.schedulerType = scheduler
82+
config.schedulerType = scheduler.asStableDiffusionScheduler()
8383
config.useDenoisedIntermediates = true
8484

8585
// Evenly distribute previews based on inference steps

0 commit comments

Comments
 (0)