tqdm auto instead

Solve issue #26

denoising_diffusion_pytorch/denoising_diffusion_pytorch.py

@@ -15,10 +15,10 @@ from torch.optim import Adam
 from torchvision import transforms, utils
 from PIL import Image
 
-from tqdm import tqdm
 from einops import rearrange, reduce
 from einops.layers.torch import Rearrange
 
+from tqdm.auto import tqdm
 from ema_pytorch import EMA
 
 # helpers functions

denoising_diffusion_pytorch/elucidated_diffusion.py

@@ -96,13 +96,30 @@ class ElucidatedDiffusion(nn.Module):
     def c_noise(self, sigma):
         return log(sigma) * 0.25
 
-    # noise distribution
+    # preconditioned network output
+    # equation (7) in the paper
 
-    def noise_distribution(self, batch_size):
-        return (self.P_mean + self.P_std * torch.randn((batch_size,), device = self.device)).exp()
+    def preconditioned_network_forward(self, noised_images, sigma, clamp = False):
+        batch, device = noised_images.shape[0], noised_images.device
 
-    def loss_weight(self, sigma):
-        return (sigma ** 2 + self.sigma_data ** 2) * (sigma * self.sigma_data) ** -2
+        if isinstance(sigma, float):
+            sigma = torch.full((batch,), sigma, device = device)
+
+        padded_sigma = rearrange(sigma, 'b -> b 1 1 1')
+
+        net_out = self.net(
+            self.c_in(padded_sigma) * noised_images,
+            self.c_noise(sigma)
+        )
+
+        out = self.c_skip(padded_sigma) * noised_images +  self.c_out(padded_sigma) * net_out
+
+        if clamp:
+            out = out.clamp(-1., 1.)
+
+        return out
+
+    # sampling
 
     # sample schedule
     # equation (5) in the paper
@@ -119,28 +136,8 @@ class ElucidatedDiffusion(nn.Module):
         sigmas = F.pad(sigmas, (0, 1), value = 0.) # last step is sigma value of 0.
         return sigmas
 
-    # preconditioned network output
-    # equation (7) in the paper
-
-    def preconditioned_network_forward(self, noised_images, sigma):
-        batch, device = noised_images.shape[0], noised_images.device
-
-        if isinstance(sigma, float):
-            sigma = torch.full((batch,), sigma, device = device)
-
-        padded_sigma = rearrange(sigma, 'b -> b 1 1 1')
-
-        net_out = self.net(
-            self.c_in(padded_sigma) * noised_images,
-            self.c_noise(sigma)
-        )
-
-        return self.c_skip(padded_sigma) * noised_images +  self.c_out(padded_sigma) * net_out
-
-    # sampling
-
     @torch.no_grad()
-    def sample(self, batch_size = 16, num_sample_steps = None):
+    def sample(self, batch_size = 16, num_sample_steps = None, clamp = True):
         num_sample_steps = default(num_sample_steps, self.num_sample_steps)
 
         shape = (batch_size, self.channels, self.image_size, self.image_size)
@@ -168,12 +165,12 @@ class ElucidatedDiffusion(nn.Module):
         for sigma, sigma_next, gamma in tqdm(sigmas_and_gammas, desc = 'sampling time step'):
             sigma, sigma_next, gamma = map(lambda t: t.item(), (sigma, sigma_next, gamma))
 
-            eps = gamma * torch.randn(shape, device = self.device)
+            eps = self.S_noise * torch.randn(shape, device = self.device) # stochastic sampling
 
             sigma_hat = sigma + gamma * sigma
             images_hat = images + sqrt(sigma_hat ** 2 - sigma ** 2) * eps
 
-            model_output = self.preconditioned_network_forward(images_hat, sigma_hat)
+            model_output = self.preconditioned_network_forward(images_hat, sigma_hat, clamp = clamp)
             denoised_over_sigma = (images_hat - model_output) / sigma_hat
 
             images_next = images_hat + (sigma_next - sigma_hat) * denoised_over_sigma
@@ -181,7 +178,7 @@ class ElucidatedDiffusion(nn.Module):
             # second order correction, if not the last timestep
 
             if sigma_next != 0:
-                model_output_next = self.preconditioned_network_forward(images_next, sigma_next)
+                model_output_next = self.preconditioned_network_forward(images_next, sigma_next, clamp = clamp)
                 denoised_prime_over_sigma = (images_next - model_output_next) / sigma_next
                 images_next = images_hat + 0.5 * (sigma_next - sigma_hat) * (denoised_over_sigma + denoised_prime_over_sigma)
 
@@ -192,6 +189,12 @@ class ElucidatedDiffusion(nn.Module):
 
     # training
 
+    def loss_weight(self, sigma):
+        return (sigma ** 2 + self.sigma_data ** 2) * (sigma * self.sigma_data) ** -2
+
+    def noise_distribution(self, batch_size):
+        return (self.P_mean + self.P_std * torch.randn((batch_size,), device = self.device)).exp()
+
     def forward(self, images):
         batch_size, c, h, w, device, image_size, channels = *images.shape, images.device, self.image_size, self.channels
 

setup.py

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