From 79e98f445707efcfd78a5dfe6351997cea49368b Mon Sep 17 00:00:00 2001
From: Pedro Cuenca <pedro@huggingface.co>
Date: Wed, 18 Jan 2023 10:54:08 +0100
Subject: [PATCH] Refactor DiffusionGlobals to GenerationContext, includes
 generation state.

---
 Diffusion-macOS/PromptView.swift  |  2 +-
 Diffusion/State.swift             |  9 ++++++++-
 Diffusion/Views/Loading.swift     |  2 +-
 Diffusion/Views/TextToImage.swift | 17 +++++------------
 4 files changed, 15 insertions(+), 15 deletions(-)

diff --git a/Diffusion-macOS/PromptView.swift b/Diffusion-macOS/PromptView.swift
index 7bcc523..3b1a904 100644
--- a/Diffusion-macOS/PromptView.swift
+++ b/Diffusion-macOS/PromptView.swift
@@ -18,7 +18,7 @@ enum PipelineState {
 }
 
 struct PromptView: View {
-    @StateObject var context = DiffusionGlobals()
+    @StateObject var context = GenerationContext()
 
     static let models = ModelInfo.MODELS
     static let modelNames = models.map { $0.modelVersion }
diff --git a/Diffusion/State.swift b/Diffusion/State.swift
index b0a389b..56c149a 100644
--- a/Diffusion/State.swift
+++ b/Diffusion/State.swift
@@ -9,8 +9,15 @@ import SwiftUI
 
 let DEFAULT_MODEL = ModelInfo.v2Base
 
-class DiffusionGlobals: ObservableObject {
+enum GenerationState {
+    case startup
+    case running(StableDiffusionProgress?)
+    case complete(String, CGImage?, TimeInterval?)
+}
+
+class GenerationContext: ObservableObject {
     @Published var pipeline: Pipeline? = nil
+    @Published var state: GenerationState = .startup
 }
 
 class Settings {
diff --git a/Diffusion/Views/Loading.swift b/Diffusion/Views/Loading.swift
index d077f40..ae9f727 100644
--- a/Diffusion/Views/Loading.swift
+++ b/Diffusion/Views/Loading.swift
@@ -12,7 +12,7 @@ import Combine
 let model = ModelInfo.v2Base
 
 struct LoadingView: View {
-    @StateObject var context = DiffusionGlobals()
+    @StateObject var context = GenerationContext()
 
     @State private var preparationPhase = "Downloading…"
     @State private var downloadProgress: Double = 0
diff --git a/Diffusion/Views/TextToImage.swift b/Diffusion/Views/TextToImage.swift
index 779b97b..7b78644 100644
--- a/Diffusion/Views/TextToImage.swift
+++ b/Diffusion/Views/TextToImage.swift
@@ -20,12 +20,6 @@ func generate(pipeline: Pipeline?, prompt: String) async -> (CGImage, TimeInterv
     return try? pipeline.generate(prompt: prompt, scheduler: scheduler, numInferenceSteps: steps, seed: seed)
 }
 
-enum GenerationState {
-    case startup
-    case running(StableDiffusionProgress?)
-    case idle(String, TimeInterval?)
-}
-
 /// 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 {
@@ -67,7 +61,6 @@ struct ShareButtons: View {
 }
 
 struct ImageWithPlaceholder: View {
-    var image: Binding<CGImage?>
     var state: Binding<GenerationState>
         
     var body: some View {
@@ -82,8 +75,8 @@ struct ImageWithPlaceholder: 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 .idle(let lastPrompt, let interval):
-            guard let theImage = image.wrappedValue else {
+        case .complete(let lastPrompt, let image, let interval):
+            guard let theImage = image else {
                 return AnyView(Image(systemName: "exclamationmark.triangle").resizable())
             }
                               
@@ -107,7 +100,7 @@ struct ImageWithPlaceholder: View {
 }
 
 struct TextToImage: View {
-    @EnvironmentObject var context: DiffusionGlobals
+    @EnvironmentObject var context: GenerationContext
 
     @State private var prompt = "Labrador in the style of Vermeer"
     @State private var image: CGImage? = nil
@@ -121,7 +114,7 @@ struct TextToImage: View {
             state = .running(nil)
             let interval: TimeInterval?
             (image, interval) = await generate(pipeline: context.pipeline, prompt: prompt) ?? (nil, nil)
-            state = .idle(prompt, interval)
+            state = .complete(prompt, image, interval)
         }
     }
     
@@ -139,7 +132,7 @@ struct TextToImage: View {
                 .padding()
                 .buttonStyle(.borderedProminent)
             }
-            ImageWithPlaceholder(image: $image, state: $state)
+            ImageWithPlaceholder(state: $state)
                 .scaledToFit()
             Spacer()
         }