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

Author: lucidrains

denoising_diffusion_pytorch/continuous_time_gaussian_diffusion.py

@@ -108,18 +108,16 @@ def p_mean_variance(self, x, time, time_next):
         # todo - derive x_start from the posterior mean and do dynamic thresholding
         # assumed that is what is going on in Imagen
 
-        batch = x.shape[0]
-        batch_time = repeat(time, ' -> b', b = batch)
-
-        pred_noise = self.denoise_fn(x, batch_time * self.cond_scale)
-
         log_snr = self.log_snr(time)
         log_snr_next = self.log_snr(time_next)
         c = -expm1(log_snr - log_snr_next)
 
         squared_alpha, squared_alpha_next = log_snr.sigmoid(), log_snr_next.sigmoid()
         squared_sigma, squared_sigma_next = (-log_snr).sigmoid(), (-log_snr_next).sigmoid()
 
+        batch_log_snr = repeat(log_snr, ' -> b', b = x.shape[0])
+        pred_noise = self.denoise_fn(x, batch_log_snr)
+
         model_mean = sqrt(squared_alpha_next / squared_alpha) * (x - c * sqrt(squared_sigma) * pred_noise)
         posterior_variance = squared_sigma_next * c
 
@@ -151,6 +149,7 @@ def p_sample_loop(self, shape):
             times_next = steps[i + 1]
             img = self.p_sample(img, times, times_next)
 
+        img.clamp_(-1., 1.)
         img = unnormalize_to_zero_to_one(img)
         return img
 

setup.py

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