unet needs to be conditioned on log(snr) in p_mean_variance for continuous time gaussian diffusion

denoising_diffusion_pytorch/continuous_time_gaussian_diffusion.py

@@ -59,6 +59,7 @@ class ContinuousTimeGaussianDiffusion(nn.Module):
         *,
         image_size,
         channels = 3,
+        cond_scale = 500,
         loss_type = 'l1',
         noise_schedule = 'linear',
         num_sample_steps = 500
@@ -75,6 +76,7 @@ class ContinuousTimeGaussianDiffusion(nn.Module):
 
         # continuous noise schedule related stuff
 
+        self.cond_scale = cond_scale # the log(snr) will be scaled by this value
         self.loss_type = loss_type
 
         if noise_schedule == 'linear':
@@ -177,7 +179,7 @@ class ContinuousTimeGaussianDiffusion(nn.Module):
 
         x, log_snr = self.q_sample(x_start = x_start, times = times, noise = noise)
 
-        model_out = self.denoise_fn(x, log_snr)
+        model_out = self.denoise_fn(x, log_snr * self.cond_scale)
         return self.loss_fn(model_out, noise)
 
     def forward(self, img, *args, **kwargs):

setup.py

@@ -3,7 +3,7 @@ from setuptools import setup, find_packages
 setup(
   name = 'denoising-diffusion-pytorch',
   packages = find_packages(),
-  version = '0.16.7',
+  version = '0.16.6',
   license='MIT',
   description = 'Denoising Diffusion Probabilistic Models - Pytorch',
   author = 'Phil Wang',