Fix the LMS pytorch regression

Author: anton-l

src/diffusers/schedulers/scheduling_lms_discrete.py

@@ -85,25 +85,26 @@ def __init__(
             )
 
         if trained_betas is not None:
-            self.betas = torch.from_numpy(trained_betas)
+            self.betas = np.asarray(trained_betas)
         if beta_schedule == "linear":
-            self.betas = torch.linspace(beta_start, beta_end, num_train_timesteps, dtype=torch.float32)
+            self.betas = np.linspace(beta_start, beta_end, num_train_timesteps, dtype=np.float32)
         elif beta_schedule == "scaled_linear":
             # this schedule is very specific to the latent diffusion model.
             self.betas = (
-                torch.linspace(beta_start**0.5, beta_end**0.5, num_train_timesteps, dtype=torch.float32) ** 2
+                np.linspace(beta_start**0.5, beta_end**0.5, num_train_timesteps, dtype=np.float32) ** 2
             )
         else:
             raise NotImplementedError(f"{beta_schedule} does is not implemented for {self.__class__}")
 
-        self.alphas = 1.0 - self.betas
-        self.alphas_cumprod = torch.cumprod(self.alphas, dim=0)
+        self.alphas = np.array(1.0 - self.betas, dtype=np.float32)
+        self.alphas_cumprod = np.cumprod(self.alphas, axis=0)
 
-        self.sigmas = ((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5
+        sigmas = ((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5
+        self.sigmas = torch.from_numpy(sigmas)
 
         # setable values
         self.num_inference_steps = None
-        self.timesteps = np.arange(0, num_train_timesteps)[::-1]  # to be consistent has to be smaller than sigmas by 1
+        self.timesteps = np.arange(0, num_train_timesteps)[::-1]
         self.derivatives = []
 
     def get_lms_coefficient(self, order, t, current_order):
@@ -146,8 +147,8 @@ def set_timesteps(self, num_inference_steps: int):
         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
 
-        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.370) < 1e-2
+        assert abs(result_sum.item() - 1006.388) < 1e-2
         assert abs(result_mean.item() - 1.31) < 1e-3