From 3c30af52eeba84e7df21ca41132aa5ee09726424 Mon Sep 17 00:00:00 2001
From: Timothy Kautz <tim@timothykautz.com>
Date: Tue, 20 Dec 2022 06:31:29 +0900
Subject: [PATCH] clamp initial timestamp

---
 swift/StableDiffusion/pipeline/AlphasCumprodCalculation.swift | 2 +-
 swift/StableDiffusion/pipeline/Encoder.swift                  | 1 -
 2 files changed, 1 insertion(+), 2 deletions(-)

diff --git a/swift/StableDiffusion/pipeline/AlphasCumprodCalculation.swift b/swift/StableDiffusion/pipeline/AlphasCumprodCalculation.swift
index 350ba3a..7875683 100644
--- a/swift/StableDiffusion/pipeline/AlphasCumprodCalculation.swift
+++ b/swift/StableDiffusion/pipeline/AlphasCumprodCalculation.swift
@@ -22,7 +22,7 @@ public struct AlphasCumprodCalculation {
         strength: Float
     ) {
         let tEnc = Int(strength * Float(steps))
-        let initTimestep = timesteps - timesteps / steps * (steps - tEnc) + 1
+        let initTimestep = min(max(0, timesteps - timesteps / steps * (steps - tEnc) + 1), timesteps - 1)
         self.sqrtAlphasCumprod = alphasCumprod[initTimestep].squareRoot()
         self.sqrtOneMinusAlphasCumprod = (1 - alphasCumprod[initTimestep]).squareRoot()
     }
diff --git a/swift/StableDiffusion/pipeline/Encoder.swift b/swift/StableDiffusion/pipeline/Encoder.swift
index 043865c..856e89a 100644
--- a/swift/StableDiffusion/pipeline/Encoder.swift
+++ b/swift/StableDiffusion/pipeline/Encoder.swift
@@ -66,7 +66,6 @@ public struct Encoder {
             else {
                 throw Error.latentOutputNotValid
             }
-            print("output.shape: \(output.shape)")
             return MLShapedArray(output)
         }