From da83fdfff1628970064f86e1c264dec5ad313091 Mon Sep 17 00:00:00 2001
From: Pedro Cuenca <pedro@huggingface.co>
Date: Fri, 23 Dec 2022 19:09:04 +0100
Subject: [PATCH] Simple display of time interval.

---
 Diffusion/Pipeline/Pipeline.swift |  7 ++++---
 Diffusion/Views/TextToImage.swift | 20 ++++++++++++++------
 2 files changed, 18 insertions(+), 9 deletions(-)

diff --git a/Diffusion/Pipeline/Pipeline.swift b/Diffusion/Pipeline/Pipeline.swift
index 7d1a7c7..caacd9e 100644
--- a/Diffusion/Pipeline/Pipeline.swift
+++ b/Diffusion/Pipeline/Pipeline.swift
@@ -29,7 +29,7 @@ class Pipeline {
         self.pipeline = pipeline
     }
     
-    func generate(prompt: String, scheduler: StableDiffusionScheduler, numInferenceSteps stepCount: Int = 50, seed: UInt32? = nil) throws -> CGImage {
+    func generate(prompt: String, scheduler: StableDiffusionScheduler, numInferenceSteps stepCount: Int = 50, seed: UInt32? = nil) throws -> (CGImage, TimeInterval) {
         let beginDate = Date()
         print("Generating...")
         let theSeed = seed ?? UInt32.random(in: 0..<UInt32.max)
@@ -43,11 +43,12 @@ class Pipeline {
             handleProgress(progress)
             return true
         }
-        print("Got images: \(images) in \(Date().timeIntervalSince(beginDate))")
+        let interval = Date().timeIntervalSince(beginDate)
+        print("Got images: \(images) in \(interval)")
         
         // unwrap the 1 image we asked for
         guard let image = images.compactMap({ $0 }).first else { throw "Generation failed" }
-        return image
+        return (image, interval)
     }
 
     func handleProgress(_ progress: StableDiffusionPipeline.Progress) {
diff --git a/Diffusion/Views/TextToImage.swift b/Diffusion/Views/TextToImage.swift
index f84a9b3..040959a 100644
--- a/Diffusion/Views/TextToImage.swift
+++ b/Diffusion/Views/TextToImage.swift
@@ -15,7 +15,7 @@ let scheduler = StableDiffusionScheduler.dpmSolverMultistepScheduler
 let steps = 25
 let seed: UInt32? = nil
 
-func generate(pipeline: Pipeline?, prompt: String) async -> CGImage? {
+func generate(pipeline: Pipeline?, prompt: String) async -> (CGImage, TimeInterval)? {
     guard let pipeline = pipeline else { return nil }
     return try? pipeline.generate(prompt: prompt, scheduler: scheduler, numInferenceSteps: steps, seed: seed)
 }
@@ -23,7 +23,7 @@ func generate(pipeline: Pipeline?, prompt: String) async -> CGImage? {
 enum GenerationState {
     case startup
     case running(StableDiffusionProgress?)
-    case idle(String)
+    case idle(String, TimeInterval?)
 }
 
 /// Presents "Share" + "Save" buttons on Mac; just "Share" on iOS/iPadOS.
@@ -82,7 +82,7 @@ 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):
+        case .idle(let lastPrompt, let interval):
             guard let theImage = image.wrappedValue else {
                 return AnyView(Image(systemName: "exclamationmark.triangle").resizable())
             }
@@ -91,7 +91,14 @@ struct ImageWithPlaceholder: View {
             return AnyView(
                 VStack {
                     imageView.resizable().clipShape(RoundedRectangle(cornerRadius: 20))
-                    ShareButtons(image: theImage, name: lastPrompt)
+                    HStack {
+                        if let interval = interval {
+                            Text(String(format:"Time: %.1fs", interval))
+                        }
+                        Spacer()
+                        ShareButtons(image: theImage, name: lastPrompt)
+                        Spacer()
+                    }
             })
         }
     }
@@ -110,8 +117,9 @@ struct TextToImage: View {
         if case .running = state { return }
         Task {
             state = .running(nil)
-            image = await generate(pipeline: context.pipeline, prompt: prompt)
-            state = .idle(prompt)
+            let interval: TimeInterval?
+            (image, interval) = await generate(pipeline: context.pipeline, prompt: prompt) ?? (nil, nil)
+            state = .idle(prompt, interval)
         }
     }