add add_noise method in LMSDiscreteScheduler, PNDMScheduler (#227)

add add_noise method in more schedulers

Author: patil-suraj

src/diffusers/schedulers/scheduling_lms_discrete.py

@@ -130,5 +130,14 @@ def add_noise(self, original_samples, noise, timesteps):
         noisy_samples = (alpha_prod**0.5) * original_samples + ((1 - alpha_prod) ** 0.5) * noise
         return noisy_samples
 
+    def add_noise(self, original_samples, noise, timesteps):
+        sqrt_alpha_prod = self.alphas_cumprod[timesteps] ** 0.5
+        sqrt_alpha_prod = self.match_shape(sqrt_alpha_prod, original_samples)
+        sqrt_one_minus_alpha_prod = (1 - self.alphas_cumprod[timesteps]) ** 0.5
+        sqrt_one_minus_alpha_prod = self.match_shape(sqrt_one_minus_alpha_prod, original_samples)
+
+        noisy_samples = sqrt_alpha_prod * original_samples + sqrt_one_minus_alpha_prod * noise
+        return noisy_samples
+
     def __len__(self):
         return self.config.num_train_timesteps

src/diffusers/schedulers/scheduling_pndm.py

@@ -250,5 +250,14 @@ def _get_prev_sample(self, sample, timestep, timestep_prev, model_output):
 
         return prev_sample
 
+    def add_noise(self, original_samples, noise, timesteps):
+        sqrt_alpha_prod = self.alphas_cumprod[timesteps] ** 0.5
+        sqrt_alpha_prod = self.match_shape(sqrt_alpha_prod, original_samples)
+        sqrt_one_minus_alpha_prod = (1 - self.alphas_cumprod[timesteps]) ** 0.5
+        sqrt_one_minus_alpha_prod = self.match_shape(sqrt_one_minus_alpha_prod, original_samples)
+
+        noisy_samples = sqrt_alpha_prod * original_samples + sqrt_one_minus_alpha_prod * noise
+        return noisy_samples
+
     def __len__(self):
         return self.config.num_train_timesteps