Pass seed back, report generation errors.

Author: pcuenca

Diffusion-macOS/ContentView.swift

@@ -66,7 +66,7 @@ struct ContentView: View {
     @StateObject var generation = GenerationContext()
 
     func toolbar() -> any View {
-        if case .complete(let prompt, let cgImage, _) = generation.state, let cgImage = cgImage {
+        if case .complete(let prompt, let cgImage, let seed, _) = generation.state, let cgImage = cgImage {
             return ShareButtons(image: cgImage, name: prompt)
         } else {
             let prompt = DEFAULT_PROMPT

Diffusion-macOS/GeneratedImageView.swift

@@ -23,7 +23,7 @@ struct GeneratedImageView: View {
             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, _):
+        case .complete(_, let image, _, _):
             guard let theImage = image else {
                 return AnyView(Image(systemName: "exclamationmark.triangle").resizable())
             }
@@ -32,6 +32,8 @@ struct GeneratedImageView: View {
                 .resizable()
                 .clipShape(RoundedRectangle(cornerRadius: 20))
             )
+        case .failed(_):
+            return AnyView(Image(systemName: "exclamationmark.triangle").resizable())
         }
     }
 }

Diffusion-macOS/StatusView.swift

@@ -18,13 +18,42 @@ struct StatusView: View {
         if case .running = generation.state { return }
         Task {
             generation.state = .running(nil)
-            let interval: TimeInterval?
-            let image: CGImage?
-            (image, interval) = await generation.generate() ?? (nil, nil)
-            generation.state = .complete(generation.positivePrompt, image, interval)
+            do {
+                let result = try await generation.generate()
+                generation.state = .complete(generation.positivePrompt, result.image, result.lastSeed, result.interval)
+            } catch {
+                generation.state = .failed(error)
+            }
         }
     }
-    
+
+    func errorWithDetails(_ message: String, error: Error) -> any View {
+        HStack {
+            Text(message)
+            Spacer()
+            Button {
+                showErrorPopover.toggle()
+            } label: {
+                Image(systemName: "info.circle")
+            }.buttonStyle(.plain)
+            .popover(isPresented: $showErrorPopover) {
+                VStack {
+                    Text(verbatim: "\(error)")
+                    .lineLimit(nil)
+                    .padding(.all, 5)
+                    Button {
+                        showErrorPopover.toggle()
+                    } label: {
+                        Text("Dismiss").frame(maxWidth: 200)
+                    }
+                    .padding(.bottom)
+                }
+                .frame(minWidth: 400, idealWidth: 400, maxWidth: 400)
+                .fixedSize()
+            }
+        }
+    }
+
     func generationStatusView() -> any View {
         switch generation.state {
         case .startup: return EmptyView()
@@ -42,7 +71,7 @@ struct StatusView: View {
                 Text("Generating \(Int(round(100*fraction)))%")
                 Spacer()
             }
-        case .complete(_, let image, let interval):
+        case .complete(_, let image, let lastSeed, let interval):
             guard let _ = image else {
                 return HStack {
                     Text("Safety checker triggered, please try a different prompt or seed")
@@ -55,9 +84,11 @@ struct StatusView: View {
                 Text(intervalString)
                 Spacer()
             }.frame(maxHeight: 25)
+        case .failed(let error):
+            return errorWithDetails("Generation error", error: error)
         }
     }
-
+    
     var body: some View {
         switch pipelineState.wrappedValue {
         case .downloading(let progress):
@@ -80,30 +111,7 @@ struct StatusView: View {
                 AnyView(generationStatusView())
             }
         case .failed(let error):
-            HStack {
-                Text("Pipeline loading error")
-                Spacer()
-                Button {
-                    showErrorPopover.toggle()
-                } label: {
-                    Image(systemName: "info.circle")
-                }.buttonStyle(.plain)
-                .popover(isPresented: $showErrorPopover) {
-                    VStack {
-                        Text(verbatim: "\(error)")
-                        .lineLimit(nil)
-                        .padding(.all, 5)
-                        Button {
-                            showErrorPopover.toggle()
-                        } label: {
-                            Text("Dismiss").frame(maxWidth: 200)
-                        }
-                        .padding(.bottom)
-                    }
-                    .frame(minWidth: 400, idealWidth: 400, maxWidth: 400)
-                    .fixedSize()
-                }
-            }
+            AnyView(errorWithDetails("Pipeline loading error", error: error))
         }
     }
 }

Diffusion/Pipeline/Pipeline.swift

@@ -14,6 +14,12 @@ import StableDiffusion
 
 typealias StableDiffusionProgress = StableDiffusionPipeline.Progress
 
+struct GenerationResult {
+    var image: CGImage?
+    var lastSeed: UInt32
+    var interval: TimeInterval?
+}
+
 class Pipeline {
     let pipeline: StableDiffusionPipeline
 
@@ -37,7 +43,7 @@ class Pipeline {
         seed: UInt32? = nil,
         guidanceScale: Float = 7.5,
         disableSafety: Bool = false
-    ) throws -> (CGImage, TimeInterval) {
+    ) throws -> GenerationResult {
         let beginDate = Date()
         print("Generating...")
         let theSeed = seed ?? UInt32.random(in: 0..<UInt32.max)
@@ -59,7 +65,7 @@ class Pipeline {
 
         // unwrap the 1 image we asked for
         guard let image = images.compactMap({ $0 }).first else { throw "Generation failed" }
-        return (image, interval)
+        return GenerationResult(image: image, lastSeed: theSeed, interval: interval)
     }
 
     func handleProgress(_ progress: StableDiffusionPipeline.Progress) {

Diffusion/State.swift

@@ -16,7 +16,8 @@ let DEFAULT_PROMPT = "Labrador in the style of Vermeer"
 enum GenerationState {
     case startup
     case running(StableDiffusionProgress?)
-    case complete(String, CGImage?, TimeInterval?)
+    case complete(String, CGImage?, UInt32, TimeInterval?)
+    case failed(Error)
 }
 
 class GenerationContext: ObservableObject {
@@ -49,10 +50,10 @@ class GenerationContext: ObservableObject {
 
     private var progressSubscriber: Cancellable?
 
-    func generate() async -> (CGImage, TimeInterval)? {
-        guard let pipeline = pipeline else { return nil }
+    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(
+        return try pipeline.generate(
             prompt: positivePrompt,
             negativePrompt: negativePrompt,
             scheduler: scheduler,

Diffusion/Views/TextToImage.swift

@@ -99,7 +99,7 @@ struct TextToImage: View {
             generation.state = .running(nil)
             let interval: TimeInterval?
             let image: CGImage?
-            (image, interval) = await generation.generate() ?? (nil, nil)
+            let result = await generation.generate()
             generation.state = .complete(generation.positivePrompt, image, interval)
         }
     }