From d9563fb5867a8ca51535283449365547daab663f Mon Sep 17 00:00:00 2001
From: Pedro Cuenca <pedro@huggingface.co>
Date: Wed, 28 Dec 2022 20:23:38 +0100
Subject: [PATCH] Store timesteps in reverse order for consistency.

---
 .../pipeline/DPMSolverMultistepScheduler.swift        |  7 +------
 swift/StableDiffusion/pipeline/Scheduler.swift        | 11 +++++------
 .../pipeline/StableDiffusionPipeline.swift            |  7 -------
 3 files changed, 6 insertions(+), 19 deletions(-)

diff --git a/swift/StableDiffusion/pipeline/DPMSolverMultistepScheduler.swift b/swift/StableDiffusion/pipeline/DPMSolverMultistepScheduler.swift
index 97c8525..1c3f261 100644
--- a/swift/StableDiffusion/pipeline/DPMSolverMultistepScheduler.swift
+++ b/swift/StableDiffusion/pipeline/DPMSolverMultistepScheduler.swift
@@ -178,10 +178,5 @@ public final class DPMSolverMultistepScheduler: Scheduler {
         }
         
         return prevSample
-    }
-    
-    /// This scheduler does not support image2image strength value.
-    public func calculateTimesteps(strength: Float?) -> [Int] {
-        timeSteps
-    }
+    }    
 }
diff --git a/swift/StableDiffusion/pipeline/Scheduler.swift b/swift/StableDiffusion/pipeline/Scheduler.swift
index d70a42d..ee649f8 100644
--- a/swift/StableDiffusion/pipeline/Scheduler.swift
+++ b/swift/StableDiffusion/pipeline/Scheduler.swift
@@ -92,12 +92,11 @@ public extension Scheduler {
 
 @available(iOS 16.2, macOS 13.1, *)
 public extension Scheduler {
-    
     func calculateTimesteps(strength: Float?) -> [Int] {
-        guard let strength else { return timeSteps.reversed() }
-        let startStep = Int(Float(inferenceStepCount) * strength)
-        let acutalTimesteps = Array(timeSteps[0..<startStep].reversed())
-        return acutalTimesteps
+        guard let strength else { return timeSteps }
+        let startStep = max(inferenceStepCount - Int(Float(inferenceStepCount) * strength), 0)
+        let actualTimesteps = Array(timeSteps[startStep...])
+        return actualTimesteps
     }
 }
 
@@ -175,7 +174,7 @@ public final class PNDMScheduler: Scheduler {
         timeSteps.append(contentsOf: forwardSteps.dropLast(1))
         timeSteps.append(timeSteps.last!)
         timeSteps.append(forwardSteps.last!)
-        // do no revers timeSteps, this is now done in `calculateTimesteps` function
+        timeSteps.reverse()
 
         self.timeSteps = timeSteps
         self.counter = 0
diff --git a/swift/StableDiffusion/pipeline/StableDiffusionPipeline.swift b/swift/StableDiffusion/pipeline/StableDiffusionPipeline.swift
index a0ba824..7e54b91 100644
--- a/swift/StableDiffusion/pipeline/StableDiffusionPipeline.swift
+++ b/swift/StableDiffusion/pipeline/StableDiffusionPipeline.swift
@@ -23,7 +23,6 @@ public struct StableDiffusionPipeline: ResourceManaging {
     
     public enum Error: String, Swift.Error {
         case startingImageProvidedWithoutEncoder
-        case schedulerNotSupportedWithImage2Image
     }
 
     /// Model to generate embeddings for tokenized input text
@@ -170,12 +169,6 @@ public struct StableDiffusionPipeline: ResourceManaging {
             guard let encoder else {
                 throw Error.startingImageProvidedWithoutEncoder
             }
-            switch schedulerType {
-            case .pndmScheduler:
-                break
-            case .dpmSolverMultistepScheduler:
-                throw Error.schedulerNotSupportedWithImage2Image
-            }
             
             let noiseTuples = generateImage2ImageLatentSamples(imageCount, stdev: 1, seed: seed)
             latents = try noiseTuples.map({