From 1f0c5b4cf196716c21f2db74f92c8750ad8f889d Mon Sep 17 00:00:00 2001
From: wfng92 <43742196+wfng92@users.noreply.github.com>
Date: Thu, 3 Nov 2022 17:13:52 +0800
Subject: [PATCH] Use array slicing to calc ddim timesteps

---
 ldm/modules/diffusionmodules/util.py | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/ldm/modules/diffusionmodules/util.py b/ldm/modules/diffusionmodules/util.py
index e93cf49a4d..c6b2eb4c07 100644
--- a/ldm/modules/diffusionmodules/util.py
+++ b/ldm/modules/diffusionmodules/util.py
@@ -65,8 +65,10 @@ def make_ddim_timesteps(
     if ddim_discr_method == 'uniform':
         c = num_ddpm_timesteps // num_ddim_timesteps
         if c < 1:
-          c = 1
-        ddim_timesteps = (np.arange(0, num_ddim_timesteps) * c).astype(int)
+            c = 1
+        
+        # remove 1 final step to prevent index out of bound error
+        ddim_timesteps = np.asarray(list(range(0, num_ddpm_timesteps, c)))[:-1]
     elif ddim_discr_method == 'quad':
         ddim_timesteps = (
             (
@@ -84,7 +86,6 @@ def make_ddim_timesteps(
     # assert ddim_timesteps.shape[0] == num_ddim_timesteps
     # add one to get the final alpha values right (the ones from first scale to data during sampling)
     steps_out = ddim_timesteps + 1
-    # steps_out = ddim_timesteps
 
     if verbose:
         print(f'Selected timesteps for ddim sampler: {steps_out}')