From 601a081737d84073c99c2dddce561ca74fa8dcbd Mon Sep 17 00:00:00 2001
From: Pedro Cuenca <pedro@huggingface.co>
Date: Fri, 23 Dec 2022 12:14:14 +0100
Subject: [PATCH] Show explicit Save button on macOS.

---
 Diffusion/Diffusion.entitlements  |  2 +-
 Diffusion/Views/TextToImage.swift | 44 +++++++++++++++++++++++++++++--
 2 files changed, 43 insertions(+), 3 deletions(-)

diff --git a/Diffusion/Diffusion.entitlements b/Diffusion/Diffusion.entitlements
index cc6a7a1..566619c 100644
--- a/Diffusion/Diffusion.entitlements
+++ b/Diffusion/Diffusion.entitlements
@@ -8,7 +8,7 @@
 	<true/>
 	<key>com.apple.security.app-sandbox</key>
 	<true/>
-	<key>com.apple.security.files.user-selected.read-only</key>
+	<key>com.apple.security.files.user-selected.read-write</key>
 	<true/>
 	<key>com.apple.security.network.client</key>
 	<true/>
diff --git a/Diffusion/Views/TextToImage.swift b/Diffusion/Views/TextToImage.swift
index 350dc10..f84a9b3 100644
--- a/Diffusion/Views/TextToImage.swift
+++ b/Diffusion/Views/TextToImage.swift
@@ -26,6 +26,46 @@ enum GenerationState {
     case idle(String)
 }
 
+/// Presents "Share" + "Save" buttons on Mac; just "Share" on iOS/iPadOS.
+/// This is because I didn't find a way for "Share" to show a Save option when running on macOS.
+struct ShareButtons: View {
+    var image: CGImage
+    var name: String
+    
+    var filename: String {
+        name.replacingOccurrences(of: " ", with: "_")
+    }
+    
+    var body: some View {
+        let imageView = Image(image, scale: 1, label: Text(name))
+
+        if (ProcessInfo.processInfo.isMacCatalystApp) {
+            HStack {
+                ShareLink(item: imageView, preview: SharePreview(name, image: imageView))
+                Button() {
+                    guard let imageData = UIImage(cgImage: image).pngData() else {
+                        return
+                    }
+                    do {
+                        let fileURL = FileManager.default.temporaryDirectory.appendingPathComponent("\(filename).png")
+                        try imageData.write(to: fileURL)
+                        let controller = UIDocumentPickerViewController(forExporting: [fileURL])
+                        
+                        let scene = UIApplication.shared.connectedScenes.first as! UIWindowScene
+                        scene.windows.first!.rootViewController!.present(controller, animated: true)
+                    } catch {
+                        print("Error creating file")
+                    }
+                } label: {
+                    Label("Save…", systemImage: "square.and.arrow.down")
+                }
+            }
+        } else {
+            ShareLink(item: imageView, preview: SharePreview(name, image: imageView))
+        }
+    }
+}
+
 struct ImageWithPlaceholder: View {
     var image: Binding<CGImage?>
     var state: Binding<GenerationState>
@@ -50,8 +90,8 @@ struct ImageWithPlaceholder: View {
             let imageView = Image(theImage, scale: 1, label: Text("generated"))
             return AnyView(
                 VStack {
-                imageView.resizable().clipShape(RoundedRectangle(cornerRadius: 20))
-                    ShareLink(item: imageView, preview: SharePreview(lastPrompt, image: imageView))
+                    imageView.resizable().clipShape(RoundedRectangle(cornerRadius: 20))
+                    ShareButtons(image: theImage, name: lastPrompt)
             })
         }
     }