AI/diffusers
1
0
Fork 0
mirror of https://github.com/huggingface/diffusers.git synced 2026-01-29 07:22:12 +03:00
Code Activity

add add_noise method in LMSDiscreteScheduler, PNDMScheduler (#227)

Browse Source 
add add_noise method in more schedulers
This commit is contained in:
Suraj Patil
2022-08-22 12:08:07 +05:30
committed by GitHub
parent 3f1861ee46
commit 5321f3e203
2 changed files with 18 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
src/diffusers/schedulers/scheduling_lms_discrete.py
View File

@@ -130,5 +130,14 @@ class LMSDiscreteScheduler(SchedulerMixin, ConfigMixin):
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

9
src/diffusers/schedulers/scheduling_pndm.py
View File

@@ -250,5 +250,14 @@ class PNDMScheduler(SchedulerMixin, ConfigMixin):
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