Preview progress images (huggingface#65)

* Render preview image per step, upgrade steps UI

* Update UI for preview options

* Reset preview image if disabled or at startup

* Update preview for ios

* Update preview help content text

* Add new option for denoised intermediates

* Cleanup edits from CR

- Also includes new Package.resolved

Author: ZachNagengast

Diffusion-macOS/ControlsView.swift

@@ -55,6 +55,7 @@ struct ControlsView: View {
     @State private var disclosedPrompt = true
     @State private var disclosedGuidance = false
     @State private var disclosedSteps = false
+    @State private var disclosedPreview = false
     @State private var disclosedSeed = false
     @State private var disclosedAdvanced = false
 
@@ -71,6 +72,7 @@ struct ControlsView: View {
     @State private var showPromptsHelp = false
     @State private var showGuidanceHelp = false
     @State private var showStepsHelp = false
+    @State private var showPreviewHelp = false
     @State private var showSeedHelp = false
     @State private var showAdvancedHelp = false
     @State private var positiveTokenCount: Int = 0
@@ -261,7 +263,7 @@ struct ControlsView: View {
                     }
 
                     DisclosureGroup(isExpanded: $disclosedSteps) {
-                        CompactSlider(value: $generation.steps, in: 0...150, step: 5) {
+                        CompactSlider(value: $generation.steps, in: 1...150, step: 1) {
                             Text("Steps")
                             Spacer()
                             Text("\(Int(generation.steps))")
@@ -285,7 +287,33 @@ struct ControlsView: View {
                             }
                         }.foregroundColor(.secondary)
                     }
-                                        
+
+                    DisclosureGroup(isExpanded: $disclosedPreview) {
+                        CompactSlider(value: $generation.previews, in: 0...25, step: 1) {
+                            Text("Previews")
+                            Spacer()
+                            Text("\(Int(generation.previews))")
+                        }.padding(.leading, 10)
+                    } label: {
+                        HStack {
+                            Label("Preview count", systemImage: "eye.square").foregroundColor(.secondary)
+                            Spacer()
+                            if disclosedPreview {
+                                Button {
+                                    showPreviewHelp.toggle()
+                                } label: {
+                                    Image(systemName: "info.circle")
+                                }
+                                .buttonStyle(.plain)
+                                .popover(isPresented: $showPreviewHelp, arrowEdge: .trailing) {
+                                    previewHelp($showPreviewHelp)
+                                }
+                            } else {
+                                Text("\(Int(generation.previews))")
+                            }
+                        }.foregroundColor(.secondary)
+                    }
+
                     DisclosureGroup(isExpanded: $disclosedSeed) {
                         let sliderLabel = generation.seed < 0 ? "Random Seed" : "Seed"
                         CompactSlider(value: $generation.seed, in: -1...Double(maxSeed), step: 1) {
@@ -312,7 +340,7 @@ struct ControlsView: View {
                             }
                         }.foregroundColor(.secondary)
                     }
-                    
+
                     if Capabilities.hasANE {
                         Divider()
                         DisclosureGroup(isExpanded: $disclosedAdvanced) {

Diffusion-macOS/GeneratedImageView.swift

@@ -10,7 +10,7 @@ import SwiftUI
 
 struct GeneratedImageView: View {
     @EnvironmentObject var generation: GenerationContext
-    
+
     var body: some View {
         switch generation.state {
         case .startup: return AnyView(Image("placeholder").resizable())
@@ -19,23 +19,34 @@ struct GeneratedImageView: View {
                 // The first time it takes a little bit before generation starts
                 return AnyView(ProgressView())
             }
+
             let step = Int(progress.step) + 1
             let fraction = Double(step) / Double(progress.stepCount)
             let label = "Step \(step) of \(progress.stepCount)"
-            return AnyView(HStack {
-                ProgressView(label, value: fraction, total: 1).padding()
-                Button {
-                    generation.cancelGeneration()
-                } label: {
-                    Image(systemName: "x.circle.fill").foregroundColor(.gray)
+
+            return AnyView(VStack {
+                Group {
+                    if let safeImage = generation.previewImage {
+                        Image(safeImage, scale: 1, label: Text("generated"))
+                            .resizable()
+                            .clipShape(RoundedRectangle(cornerRadius: 20))                        
+                    }
+                }
+                HStack {
+                    ProgressView(label, value: fraction, total: 1).padding()
+                    Button {
+                        generation.cancelGeneration()
+                    } label: {
+                        Image(systemName: "x.circle.fill").foregroundColor(.gray)
+                    }
+                    .buttonStyle(.plain)
                 }
-                .buttonStyle(.plain)
             })
         case .complete(_, let image, _, _):
             guard let theImage = image else {
                 return AnyView(Image(systemName: "exclamationmark.triangle").resizable())
             }
-                              
+            
             return AnyView(
                     Image(theImage, scale: 1, label: Text("generated"))
                     .resizable()

Diffusion-macOS/HelpContent.swift

@@ -110,6 +110,25 @@ func stepsHelp(_ showing: Binding<Bool>) -> some View {
     return helpContent(title: "Inference Steps", description: description, showing: showing)
 }
 
+func previewHelp(_ showing: Binding<Bool>) -> some View {
+    let description =
+         """
+         This number controls how many previews to display throughout the image generation process.
+
+         Using more previews can be useful if you want more visibility into how \
+         generation is progressing.
+
+         However, computing each preview takes some time and can slow down \
+         generation. If the process is too slow you can reduce the preview count, \
+         which will result in less visibility of intermediate steps during generation.
+
+         You can try different values to see what works best for your hardware.
+
+         For the absolute fastest generation times, use 0 previews.
+         """
+    return helpContent(title: "Preview Count", description: description, showing: showing)
+}
+
 func seedHelp(_ showing: Binding<Bool>) -> some View {
     let description =
          """

Diffusion.xcodeproj/project.xcworkspace/xcshareddata/swiftpm/Package.resolved

@@ -6,7 +6,7 @@
       "location" : "https://github.com/buh/CompactSlider.git",
       "state" : {
         "branch" : "main",
-        "revision" : "3cb37fb7385913835b6844c6af2680c64000dcd2"
+        "revision" : "6d591a76caecd583ad69fbcd06f2fb83135318c0"
       }
     },
     {
@@ -15,7 +15,7 @@
       "location" : "https://github.com/apple/ml-stable-diffusion",
       "state" : {
         "branch" : "main",
-        "revision" : "48f07f24891155a14c51dd835bba7371bdf32d0e"
+        "revision" : "b61c9aea05370d4bc06fce2dc00a002b21f13da5"
       }
     },
     {

Diffusion/Common/Pipeline/Pipeline.swift

@@ -12,7 +12,24 @@ import Combine
 
 import StableDiffusion
 
-typealias StableDiffusionProgress = StableDiffusionPipeline.Progress
+struct StableDiffusionProgress {
+    var progress: StableDiffusionPipeline.Progress
+
+    var step: Int { progress.step }
+    var stepCount: Int { progress.stepCount }
+
+    var currentImages: [CGImage?]
+
+    init(progress: StableDiffusionPipeline.Progress, previewIndices: [Bool]) {
+        self.progress = progress
+        self.currentImages = [nil]
+
+        // Since currentImages is a computed property, only access the preview image if necessary
+        if progress.step < previewIndices.count, previewIndices[progress.step] {
+            self.currentImages = progress.currentImages
+        }
+    }
+}
 
 struct GenerationResult {
     var image: CGImage?
@@ -45,6 +62,7 @@ class Pipeline {
         negativePrompt: String = "",
         scheduler: StableDiffusionScheduler,
         numInferenceSteps stepCount: Int = 50,
+        numPreviews previewCount: Int = 5,
         seed: UInt32? = nil,
         guidanceScale: Float = 7.5,
         disableSafety: Bool = false
@@ -63,10 +81,16 @@ class Pipeline {
         config.guidanceScale = guidanceScale
         config.disableSafety = disableSafety
         config.schedulerType = scheduler
-        
+        config.useDenoisedIntermediates = true
+
+        // Evenly distribute previews based on inference steps
+        let previewIndices = previewIndices(stepCount, previewCount)
+
         let images = try pipeline.generateImages(configuration: config) { progress in
             sampleTimer.stop()
-            handleProgress(progress, sampleTimer: sampleTimer)
+            handleProgress(StableDiffusionProgress(progress: progress,
+                                                   previewIndices: previewIndices),
+                           sampleTimer: sampleTimer)
             if progress.stepCount != progress.step {
                 sampleTimer.start()
             }
@@ -80,7 +104,7 @@ class Pipeline {
         return GenerationResult(image: image, lastSeed: theSeed, interval: interval, userCanceled: canceled, itsPerSecond: 1.0/sampleTimer.median)
     }
 
-    func handleProgress(_ progress: StableDiffusionPipeline.Progress, sampleTimer: SampleTimer) {
+    func handleProgress(_ progress: StableDiffusionProgress, sampleTimer: SampleTimer) {
         self.progress = progress
     }
 

Diffusion/Common/State.swift

@@ -35,6 +35,7 @@ class GenerationContext: ObservableObject {
                     .receive(on: DispatchQueue.main)
                     .sink { progress in
                         guard let progress = progress else { return }
+                        self.updatePreviewIfNeeded(progress)
                         self.state = .running(progress)
                     }
             }
@@ -50,12 +51,24 @@ class GenerationContext: ObservableObject {
     @Published var numImages = 1.0
     @Published var seed = -1.0
     @Published var guidanceScale = 7.5
+    @Published var previews = 5.0
     @Published var disableSafety = false
-    
+    @Published var previewImage: CGImage? = nil
+
     @Published var computeUnits: ComputeUnits = Settings.shared.userSelectedComputeUnits ?? ModelInfo.defaultComputeUnits
 
     private var progressSubscriber: Cancellable?
 
+    private func updatePreviewIfNeeded(_ progress: StableDiffusionProgress) {
+        if previews == 0 || progress.step == 0 {
+            previewImage = nil
+        }
+
+        if previews > 0, let newImage = progress.currentImages.first, newImage != nil {
+            previewImage = newImage
+        }
+    }
+
     func generate() async throws -> GenerationResult {
         guard let pipeline = pipeline else { throw "No pipeline" }
         let seed = self.seed >= 0 ? UInt32(self.seed) : nil
@@ -64,6 +77,7 @@ class GenerationContext: ObservableObject {
             negativePrompt: negativePrompt,
             scheduler: scheduler,
             numInferenceSteps: Int(steps),
+            numPreviews: Int(previews),
             seed: seed,
             guidanceScale: Float(guidanceScale),
             disableSafety: disableSafety

Diffusion/Common/Utils.swift

@@ -15,3 +15,31 @@ extension Double {
         return String(format: "\(format)", self)
     }
 }
+
+/// Returns an array of booleans that indicates at which steps a preview should be generated.
+///
+/// - Parameters:
+///   - numInferenceSteps: The total number of inference steps.
+///   - numPreviews: The desired number of previews.
+///
+/// - Returns: An array of booleans of size `numInferenceSteps`, where `true` values represent steps at which a preview should be made.
+func previewIndices(_ numInferenceSteps: Int, _ numPreviews: Int) -> [Bool] {
+    // Ensure valid parameters
+    guard numInferenceSteps > 0, numPreviews > 0 else {
+        return [Bool](repeating: false, count: numInferenceSteps)
+    }
+
+    // Compute the ideal (floating-point) step size, which represents the average number of steps between previews
+    let idealStep = Double(numInferenceSteps) / Double(numPreviews)
+
+    // Compute the actual steps at which previews should be made. For each preview, we multiply the ideal step size by the preview number, and round to the nearest integer.
+    // The result is converted to a `Set` for fast membership tests.
+    let previewIndices: Set<Int> = Set((0..<numPreviews).map { previewIndex in
+        return Int(round(Double(previewIndex) * idealStep))
+    })
+    
+    // Construct an array of booleans where each value indicates whether or not a preview should be made at that step.
+    let previewArray = (0..<numInferenceSteps).map { previewIndices.contains($0) }
+
+    return previewArray
+}

Diffusion/Views/TextToImage.swift

@@ -52,6 +52,7 @@ struct ShareButtons: View {
 }
 
 struct ImageWithPlaceholder: View {
+    @EnvironmentObject var generation: GenerationContext
     var state: Binding<GenerationState>
 
     var body: some View {
@@ -62,10 +63,20 @@ struct ImageWithPlaceholder: View {
                 // The first time it takes a little bit before generation starts
                 return AnyView(ProgressView())
             }
+
             let step = Int(progress.step) + 1
             let fraction = Double(step) / Double(progress.stepCount)
             let label = "Step \(step) of \(progress.stepCount)"
-            return AnyView(ProgressView(label, value: fraction, total: 1).padding())
+            return AnyView(VStack {
+                Group {
+                    if let safeImage = generation.previewImage {
+                        Image(safeImage, scale: 1, label: Text("generated"))
+                            .resizable()
+                            .clipShape(RoundedRectangle(cornerRadius: 20))
+                    }
+                }
+                ProgressView(label, value: fraction, total: 1).padding()
+            })
         case .complete(let lastPrompt, let image, _, let interval):
             guard let theImage = image else {
                 return AnyView(Image(systemName: "exclamationmark.triangle").resizable())
@@ -125,5 +136,6 @@ struct TextToImage: View {
             Spacer()
         }
         .padding()
+        .environmentObject(generation)
     }
 }