timesteps needs to be discrete

Author: kashif

examples/community/clip_guided_stable_diffusion.py

@@ -274,7 +274,7 @@ def __call__(
                 # the model input needs to be scaled to match the continuous ODE formulation in K-LMS
                 latent_model_input = latent_model_input / ((sigma**2 + 1) ** 0.5)
 
-            # # predict the noise residual
+            # predict the noise residual
             noise_pred = self.unet(latent_model_input, t, encoder_hidden_states=text_embeddings).sample
 
             # perform classifier free guidance

src/diffusers/schedulers/scheduling_lms_discrete.py

@@ -96,7 +96,7 @@ def __init__(
 
         # setable values
         self.num_inference_steps = None
-        self.timesteps = np.arange(0, num_train_timesteps)[::-1].copy()
+        self.timesteps = np.arange(0, num_train_timesteps)[::-1]  # to be consistent has to be smaller than sigmas by 1
         self.derivatives = []
 
     def get_lms_coefficient(self, order, t, current_order):
@@ -130,16 +130,17 @@ def set_timesteps(self, num_inference_steps: int):
                 the number of diffusion steps used when generating samples with a pre-trained model.
         """
         self.num_inference_steps = num_inference_steps
-        self.timesteps = np.linspace(self.config.num_train_timesteps - 1, 0, num_inference_steps, dtype=float)
+        timesteps = np.linspace(self.config.num_train_timesteps - 1, 0, num_inference_steps, dtype=float)
 
-        low_idx = np.floor(self.timesteps).astype(int)
-        high_idx = np.ceil(self.timesteps).astype(int)
-        frac = np.mod(self.timesteps, 1.0)
+        low_idx = np.floor(timesteps).astype(int)
+        high_idx = np.ceil(timesteps).astype(int)
+        frac = np.mod(timesteps, 1.0)
         sigmas = np.array(((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5)
         sigmas = (1 - frac) * sigmas[low_idx] + frac * sigmas[high_idx]
         sigmas = np.concatenate([sigmas, [0.0]]).astype(np.float32)
         self.sigmas = torch.from_numpy(sigmas)
 
+        self.timesteps = timesteps.astype(int)
         self.derivatives = []
 
     def step(

tests/test_scheduler.py

@@ -876,5 +876,5 @@ def test_full_loop_no_noise(self):
         result_sum = torch.sum(torch.abs(sample))
         result_mean = torch.mean(torch.abs(sample))
 
-        assert abs(result_sum.item() - 1006.388) < 1e-2
+        assert abs(result_sum.item() - 1006.370) < 1e-2
         assert abs(result_mean.item() - 1.31) < 1e-3