get working version of gaussian diffusion with continuous time (only beta linear schedule for now, but will eventually contain alpha cosine schedule as well as parameterized, learned monotonic MLP)

denoising_diffusion_pytorch/continuous_time_gaussian_diffusion.py

@@ -127,18 +127,12 @@ class ContinuousTimeGaussianDiffusion(nn.Module):
     # sampling related functions
 
     @torch.no_grad()
-    def p_sample(self, x, time, time_next, eps = 2e-4):
+    def p_sample(self, x, time, time_next):
         batch, *_, device = *x.shape, x.device
 
         model_mean, model_variance = self.p_mean_variance(x = x, time = time, time_next = time_next)
         noise = torch.randn_like(x)
-
+        return model_mean + sqrt(model_variance) * noise
-        # no noise when time is below some epsilon
-        # not sure how important this is
-
-        time = repeat(time, ' -> b', b = batch)
-        nonzero_mask = (1 - (time < eps).float()).reshape(batch, *((1,) * (len(x.shape) - 1)))
-        return model_mean + nonzero_mask * sqrt(model_variance) * noise
 
     @torch.no_grad()
     def p_sample_loop(self, shape):

setup.py

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