From d1c63dafbc476ae903dddbe5ee82b3b1ab55e9f4 Mon Sep 17 00:00:00 2001
From: Pedro Cuenca <pedro@huggingface.co>
Date: Tue, 17 Jan 2023 14:11:07 +0100
Subject: [PATCH] Download: display progress, hide Generate button.

---
 Diffusion-macOS/PromptView.swift | 34 ++++++++++++++------------------
 Diffusion-macOS/StatusView.swift | 24 ++++++++++++++++++++--
 2 files changed, 37 insertions(+), 21 deletions(-)

diff --git a/Diffusion-macOS/PromptView.swift b/Diffusion-macOS/PromptView.swift
index 02a5d73..79cee53 100644
--- a/Diffusion-macOS/PromptView.swift
+++ b/Diffusion-macOS/PromptView.swift
@@ -9,6 +9,13 @@ import Combine
 import SwiftUI
 import CompactSlider
 
+enum PipelineState {
+    case downloading(Double)
+    case uncompressing
+    case loading
+    case ready
+    case failed(Error)
+}
 
 struct PromptView: View {
     @StateObject var context = DiffusionGlobals()
@@ -24,9 +31,8 @@ struct PromptView: View {
     @State private var seed = 386.0
     
     // TODO: refactor download with similar code in Loading.swift (iOS)
-    @State private var preparationPhase = "Downloading…"
-    @State private var downloadProgress: Double = 0
     @State private var stateSubscriber: Cancellable?
+    @State private var pipelineState: PipelineState = .downloading(0)
 
     func modelDidChange(model: ModelInfo) {
         Task.init {
@@ -35,14 +41,11 @@ struct PromptView: View {
                 DispatchQueue.main.async {
                     switch state {
                     case .downloading(let progress):
-                        preparationPhase = "Downloading"
-                        downloadProgress = progress
+                        pipelineState = .downloading(progress)
                     case .uncompressing:
-                        preparationPhase = "Uncompressing"
-                        downloadProgress = 1
+                        pipelineState = .uncompressing
                     case .readyOnDisk:
-                        preparationPhase = "Loading"
-                        downloadProgress = 1
+                        pipelineState = .loading
                     default:
                         break
                     }
@@ -50,9 +53,10 @@ struct PromptView: View {
             }
             do {
                 context.pipeline = try await loader.prepare()
+                pipelineState = .ready
             } catch {
-                // TODO: expose to user
                 print("Could not load model, error: \(error)")
+                pipelineState = .failed(error)
             }
         }
     }
@@ -129,17 +133,9 @@ struct PromptView: View {
                         Label("Random Seed", systemImage: "leaf").foregroundColor(.secondary)
                     }
                 }
-                
-            }
-            Button {
-                // Generate image here
-            } label: {
-                Text("Generate")
-                    .frame(maxWidth: .infinity)
-                    .frame(height: 50)
             }
-            .buttonStyle(.borderedProminent)
-//            StatusView()
+            
+            StatusView(pipelineState: $pipelineState)
         }
         .padding()
         .onAppear {
diff --git a/Diffusion-macOS/StatusView.swift b/Diffusion-macOS/StatusView.swift
index 1eac324..f753d0e 100644
--- a/Diffusion-macOS/StatusView.swift
+++ b/Diffusion-macOS/StatusView.swift
@@ -8,13 +8,33 @@
 import SwiftUI
 
 struct StatusView: View {
+    var pipelineState: Binding<PipelineState>
+    
     var body: some View {
-        Text(/*@START_MENU_TOKEN@*/"Hello, World!"/*@END_MENU_TOKEN@*/)
+        switch pipelineState.wrappedValue {
+        case .downloading(let progress):
+            ProgressView("Downloading…", value: progress*100, total: 110).padding()
+        case .uncompressing:
+            ProgressView("Uncompressing…", value: 100, total: 110).padding()
+        case .loading:
+            ProgressView("Loading…", value: 105, total: 110).padding()
+        case .ready:
+            Button {
+                // Generate image here
+            } label: {
+                Text("Generate")
+                    .frame(maxWidth: .infinity)
+                    .frame(height: 50)
+            }
+            .buttonStyle(.borderedProminent)
+        case .failed:
+            Text("Pipeline loading error")
+        }
     }
 }
 
 struct StatusView_Previews: PreviewProvider {
     static var previews: some View {
-        StatusView()
+        StatusView(pipelineState: .constant(.downloading(0.2)))
     }
 }