[BugFix]: Fixed add_noise in LMSDiscreteScheduler (huggingface#253)

* Fixed add_noise in LMSDiscreteScheduler * Linting * Update src/diffusers/schedulers/scheduling_lms_discrete.py Co-authored-by: Anton Lozhkov <aglozhkov@gmail.com> Co-authored-by: Anton Lozhkov <aglozhkov@gmail.com>
harishanand95 · Aug 29, 2022 · da7d4cf · da7d4cf
1 parent 9e1b1ca
commit da7d4cf
Showing 1 changed file with 2 additions and 4 deletions.
diff --git a/src/diffusers/schedulers/scheduling_lms_discrete.py b/src/diffusers/schedulers/scheduling_lms_discrete.py
@@ -124,10 +124,8 @@ def step(
         return {"prev_sample": prev_sample}
 
     def add_noise(self, original_samples, noise, timesteps):
-        alpha_prod = self.alphas_cumprod[timesteps]
-        alpha_prod = self.match_shape(alpha_prod, original_samples)
-
-        noisy_samples = (alpha_prod**0.5) * original_samples + ((1 - alpha_prod) ** 0.5) * noise
+        sigmas = self.match_shape(self.sigmas, noise)
+        noisy_samples = original_samples + noise * sigmas[timesteps]
         return noisy_samples
 
     def __len__(self):