From c4991f576f901678aa4ca053058aff33ec5db7be Mon Sep 17 00:00:00 2001 From: Phil Wang Date: Wed, 8 Jun 2022 11:18:54 -0700 Subject: [PATCH] allow for configuring the hidden dimension of the monotonic mlp parameterizing the noise schedule --- .../continuous_time_gaussian_diffusion.py | 10 ++++++---- setup.py | 2 +- 2 files changed, 7 insertions(+), 5 deletions(-) diff --git a/denoising_diffusion_pytorch/continuous_time_gaussian_diffusion.py b/denoising_diffusion_pytorch/continuous_time_gaussian_diffusion.py index b4fa37c..0cc5a99 100644 --- a/denoising_diffusion_pytorch/continuous_time_gaussian_diffusion.py +++ b/denoising_diffusion_pytorch/continuous_time_gaussian_diffusion.py @@ -111,7 +111,8 @@ class ContinuousTimeGaussianDiffusion(nn.Module): loss_type = 'l1', noise_schedule = 'linear', num_sample_steps = 500, - clip_sample_after_noise = False + clip_after_noising_during_sampling = False, + learned_schedule_net_hidden_dim = 1024 ): super().__init__() assert not denoise_fn.sinusoidal_cond_mlp @@ -134,7 +135,8 @@ class ContinuousTimeGaussianDiffusion(nn.Module): self.log_snr = learned_noise_schedule( log_snr_max = log_snr_max, - log_snr_min = log_snr_min + log_snr_min = log_snr_min, + hidden_dim = learned_schedule_net_hidden_dim ) else: raise ValueError(f'unknown noise schedule {noise_schedule}') @@ -145,7 +147,7 @@ class ContinuousTimeGaussianDiffusion(nn.Module): # clipping related hyperparameters - self.clip_sample_after_noise = clip_sample_after_noise + self.clip_after_noising_during_sampling = clip_after_noising_during_sampling @property def device(self): @@ -208,7 +210,7 @@ class ContinuousTimeGaussianDiffusion(nn.Module): times_next = steps[i + 1] img = self.p_sample(img, times, times_next) - if self.clip_sample_after_noise: + if self.clip_after_noising_during_sampling: # clip after noise is added. perhaps this is sufficient for Imagen dynamic thresholding? img.clamp_(-1., 1.) diff --git a/setup.py b/setup.py index 46a2711..da27954 100644 --- a/setup.py +++ b/setup.py @@ -3,7 +3,7 @@ from setuptools import setup, find_packages setup( name = 'denoising-diffusion-pytorch', packages = find_packages(), - version = '0.17.1', + version = '0.17.2', license='MIT', description = 'Denoising Diffusion Probabilistic Models - Pytorch', author = 'Phil Wang',