0.27.6

Update step before training checkpoint

README.md

@@ -10,6 +10,8 @@ Youtube AI Educators - <a href="https://www.youtube.com/watch?v=W-O7AZNzbzQ">Yan
 
 <a href="https://huggingface.co/blog/annotated-diffusion">Annotated code</a> by Research Scientists / Engineers from <a href="https://huggingface.co/">🤗 Huggingface</a>
 
+Update: Turns out none of the technicalities really matters at all | <a href="https://arxiv.org/abs/2208.09392">"Cold Diffusion" paper</a>
+
 <img src="./images/sample.png" width="500px"><img>
 
 [![PyPI version](https://badge.fury.io/py/denoising-diffusion-pytorch.svg)](https://badge.fury.io/py/denoising-diffusion-pytorch)
@@ -181,3 +183,13 @@ $ accelerate launch train.py
     primaryClass = {cs.CV}
 }
 ```
+
+```bibtex
+@article{Qiao2019WeightS,
+    title   = {Weight Standardization},
+    author  = {Siyuan Qiao and Huiyu Wang and Chenxi Liu and Wei Shen and Alan Loddon Yuille},
+    journal = {ArXiv},
+    year    = {2019},
+    volume  = {abs/1903.10520}
+}
+```

denoising_diffusion_pytorch/denoising_diffusion_pytorch.py

@@ -88,6 +88,21 @@ def Upsample(dim, dim_out = None):
 def Downsample(dim, dim_out = None):
     return nn.Conv2d(dim, default(dim_out, dim), 4, 2, 1)
 
+class WeightStandardizedConv2d(nn.Conv2d):
+    """
+    https://arxiv.org/abs/1903.10520
+    weight standardization purportedly works synergistically with group normalization
+    """
+    def forward(self, x):
+        eps = 1e-5 if x.dtype == torch.float32 else 1e-3
+
+        weight = self.weight
+        mean = reduce(weight, 'o ... -> o 1 1 1', 'mean')
+        var = reduce(weight, 'o ... -> o 1 1 1', partial(torch.var, unbiased = False))
+        normalized_weight = (weight - mean) * (var + eps).rsqrt()
+
+        return F.conv2d(x, normalized_weight, self.bias, self.stride, self.padding, self.dilation, self.groups)
+
 class LayerNorm(nn.Module):
     def __init__(self, dim):
         super().__init__()
@@ -147,7 +162,7 @@ class LearnedSinusoidalPosEmb(nn.Module):
 class Block(nn.Module):
     def __init__(self, dim, dim_out, groups = 8):
         super().__init__()
-        self.proj = nn.Conv2d(dim, dim_out, 3, padding = 1)
+        self.proj = WeightStandardizedConv2d(dim, dim_out, 3, padding = 1)
         self.norm = nn.GroupNorm(groups, dim_out)
         self.act = nn.SiLU()
 
@@ -219,7 +234,7 @@ class LinearAttention(nn.Module):
         return self.to_out(out)
 
 class Attention(nn.Module):
-    def __init__(self, dim, heads = 4, dim_head = 32, scale = 16):
+    def __init__(self, dim, heads = 4, dim_head = 32, scale = 10):
         super().__init__()
         self.scale = scale
         self.heads = heads
@@ -236,7 +251,6 @@ class Attention(nn.Module):
 
         sim = einsum('b h d i, b h d j -> b h i j', q, k) * self.scale
         attn = sim.softmax(dim = -1)
-
         out = einsum('b h i j, b h d j -> b h i d', attn, v)
         out = rearrange(out, 'b h (x y) d -> b (h d) x y', x = h, y = w)
         return self.to_out(out)
@@ -413,6 +427,7 @@ class GaussianDiffusion(nn.Module):
     ):
         super().__init__()
         assert not (type(self) == GaussianDiffusion and model.channels != model.out_dim)
+        assert not model.learned_sinusoidal_cond
 
         self.model = model
         self.channels = self.model.channels
@@ -542,7 +557,7 @@ class GaussianDiffusion(nn.Module):
 
         x_start = None
 
-        for t in tqdm(reversed(range(0, self.num_timesteps)), desc = 'sampling loop time step'):
+        for t in tqdm(reversed(range(0, self.num_timesteps)), desc = 'sampling loop time step', total = self.num_timesteps):
             self_cond = x_start if self.self_condition else None
             img, x_start = self.p_sample(img, t, self_cond)
 
@@ -555,15 +570,15 @@ class GaussianDiffusion(nn.Module):
 
         times = torch.linspace(0., total_timesteps, steps = sampling_timesteps + 2)[:-1]
         times = list(reversed(times.int().tolist()))
-        time_pairs = list(zip(times[:-1], times[1:]))
+        time_pairs = list(filter(lambda a: a[0] > a[1], zip(times[:-1], times[1:])))
 
         img = torch.randn(shape, device = device)
 
         x_start = None
 
         for time, time_next in tqdm(time_pairs, desc = 'sampling loop time step'):
-            alpha = self.alphas_cumprod_prev[time]
-            alpha_next = self.alphas_cumprod_prev[time_next]
+            alpha = self.alphas_cumprod[time]
+            alpha_next = self.alphas_cumprod[time_next]
 
             time_cond = torch.full((batch,), time, device = device, dtype = torch.long)
 
@@ -599,11 +614,11 @@ class GaussianDiffusion(nn.Module):
 
         assert x1.shape == x2.shape
 
-        t_batched = torch.stack([torch.tensor(t, device=device)] * b)
-        xt1, xt2 = map(lambda x: self.q_sample(x, t=t_batched), (x1, x2))
+        t_batched = torch.stack([torch.tensor(t, device = device)] * b)
+        xt1, xt2 = map(lambda x: self.q_sample(x, t = t_batched), (x1, x2))
 
         img = (1 - lam) * xt1 + lam * xt2
-        for i in tqdm(reversed(range(0, t)), desc='interpolation sample time step', total=t):
+        for i in tqdm(reversed(range(0, t)), desc = 'interpolation sample time step', total = t):
             img = self.p_sample(img, torch.full((b,), i, device=device, dtype=torch.long))
 
         return img
@@ -831,6 +846,7 @@ class Trainer(object):
 
                 accelerator.wait_for_everyone()
 
+                self.step += 1
                 if accelerator.is_main_process:
                     self.ema.to(device)
                     self.ema.update()
@@ -847,7 +863,6 @@ class Trainer(object):
                         utils.save_image(all_images, str(self.results_folder / f'sample-{milestone}.png'), nrow = int(math.sqrt(self.num_samples)))
                         self.save(milestone)
 
-                self.step += 1
                 pbar.update(1)
 
         accelerator.print('training complete')

denoising_diffusion_pytorch/elucidated_diffusion.py

@@ -1,4 +1,5 @@
 from math import sqrt
+from random import random
 import torch
 from torch import nn, einsum
 import torch.nn.functional as F
@@ -52,7 +53,7 @@ class ElucidatedDiffusion(nn.Module):
     ):
         super().__init__()
         assert net.learned_sinusoidal_cond
-        assert not net.self_condition, 'not supported yet'
+        self.self_condition = net.self_condition
 
         self.net = net
 
@@ -100,7 +101,7 @@ class ElucidatedDiffusion(nn.Module):
     # preconditioned network output
     # equation (7) in the paper
 
-    def preconditioned_network_forward(self, noised_images, sigma, clamp = False):
+    def preconditioned_network_forward(self, noised_images, sigma, self_cond = None, clamp = False):
         batch, device = noised_images.shape[0], noised_images.device
 
         if isinstance(sigma, float):
@@ -110,7 +111,8 @@ class ElucidatedDiffusion(nn.Module):
 
         net_out = self.net(
             self.c_in(padded_sigma) * noised_images,
-            self.c_noise(sigma)
+            self.c_noise(sigma),
+            self_cond
         )
 
         out = self.c_skip(padded_sigma) * noised_images +  self.c_out(padded_sigma) * net_out
@@ -161,6 +163,10 @@ class ElucidatedDiffusion(nn.Module):
 
         images = init_sigma * torch.randn(shape, device = self.device)
 
+        # for self conditioning
+
+        x_start = None
+
         # gradually denoise
 
         for sigma, sigma_next, gamma in tqdm(sigmas_and_gammas, desc = 'sampling time step'):
@@ -171,7 +177,9 @@ class ElucidatedDiffusion(nn.Module):
             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, clamp = clamp)
+            self_cond = x_start if self.self_condition else None
+
+            model_output = self.preconditioned_network_forward(images_hat, sigma_hat, self_cond, clamp = clamp)
             denoised_over_sigma = (images_hat - model_output) / sigma_hat
 
             images_next = images_hat + (sigma_next - sigma_hat) * denoised_over_sigma
@@ -179,11 +187,14 @@ 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, clamp = clamp)
+                self_cond = model_output if self.self_condition else None
+
+                model_output_next = self.preconditioned_network_forward(images_next, sigma_next, self_cond, 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)
 
             images = images_next
+            x_start = model_output
 
         images = images.clamp(-1., 1.)
         return unnormalize_to_zero_to_one(images)
@@ -211,7 +222,15 @@ class ElucidatedDiffusion(nn.Module):
 
         noised_images = images + padded_sigmas * noise  # alphas are 1. in the paper
 
-        denoised = self.preconditioned_network_forward(noised_images, sigmas)
+        self_cond = None
+
+        if self.self_condition and random() < 0.5:
+            # from hinton's group's bit diffusion paper
+            with torch.no_grad():
+                self_cond = self.preconditioned_network_forward(noised_images, sigmas)
+                self_cond.detach_()
+
+        denoised = self.preconditioned_network_forward(noised_images, sigmas, self_cond)
 
         losses = F.mse_loss(denoised, images, reduction = 'none')
         losses = reduce(losses, 'b ... -> b', 'mean')

setup.py

@@ -3,7 +3,7 @@ from setuptools import setup, find_packages
 setup(
   name = 'denoising-diffusion-pytorch',
   packages = find_packages(),
-  version = '0.27.0',
+  version = '0.27.6',
   license='MIT',
   description = 'Denoising Diffusion Probabilistic Models - Pytorch',
   author = 'Phil Wang',
@@ -30,4 +30,4 @@ setup(
     'License :: OSI Approved :: MIT License',
     'Programming Language :: Python :: 3.6',
   ],
-)
+)