fix a missing residual needed at the top most resolution in the unet

README.md

@@ -6,6 +6,8 @@ Implementation of <a href="https://arxiv.org/abs/2006.11239">Denoising Diffusion
 
 This implementation was transcribed from the official Tensorflow version <a href="https://github.com/hojonathanho/diffusion">here</a>
 
+Youtube AI Educators - <a href="https://www.youtube.com/watch?v=W-O7AZNzbzQ">Yannic Kilcher</a> | <a href="https://www.youtube.com/watch?v=344w5h24-h8">AI Coffeebreak with Letitia</a> | <a href="https://www.youtube.com/watch?v=HoKDTa5jHvg">Outlier</a>
+
 <a href="https://huggingface.co/blog/annotated-diffusion">Annotated code</a> by Research Scientists / Engineers from <a href="https://huggingface.co/">🤗 Huggingface</a>
 
 <img src="./sample.png" width="500px"><img>

denoising_diffusion_pytorch/continuous_time_gaussian_diffusion.py

@@ -5,7 +5,7 @@ import torch.nn.functional as F
 from torch.special import expm1
 
 from tqdm import tqdm
-from einops import rearrange, repeat
+from einops import rearrange, repeat, reduce
 from einops.layers.torch import Rearrange
 
 # helpers
@@ -268,7 +268,7 @@ class ContinuousTimeGaussianDiffusion(nn.Module):
         model_out = self.denoise_fn(x, log_snr)
 
         losses = self.loss_fn(model_out, noise, reduction = 'none')
-        losses = losses.mean(dim = tuple(range(1, losses.ndim)))
+        losses = reduce(losses, 'b ... -> b', 'mean')
 
         if self.p2_loss_weight_gamma >= 0:
             # following eq 8. in https://arxiv.org/abs/2204.00227

denoising_diffusion_pytorch/denoising_diffusion_pytorch.py

@@ -7,6 +7,7 @@ from inspect import isfunction
 from functools import partial
 
 from torch.utils import data
+from multiprocessing import cpu_count
 from torch.cuda.amp import autocast, GradScaler
 
 from pathlib import Path
@@ -15,7 +16,7 @@ from torchvision import transforms, utils
 from PIL import Image
 
 from tqdm import tqdm
-from einops import rearrange
+from einops import rearrange, reduce
 from einops.layers.torch import Rearrange
 
 # helpers functions
@@ -301,7 +302,7 @@ class Unet(nn.Module):
         self.out_dim = default(out_dim, default_out_dim)
 
         self.final_conv = nn.Sequential(
-            block_klass(dim, dim),
+            block_klass(dim * 2, dim),
             nn.Conv2d(dim, self.out_dim, 1)
         )
 
@@ -323,12 +324,13 @@ class Unet(nn.Module):
         x = self.mid_block2(x, t)
 
         for block1, block2, attn, upsample in self.ups:
-            x = torch.cat((x, h.pop()), dim=1)
+            x = torch.cat((x, h.pop()), dim = 1)
             x = block1(x, t)
             x = block2(x, t)
             x = attn(x)
             x = upsample(x)
 
+        x = torch.cat((x, h.pop()), dim = 1)
         return self.final_conv(x)
 
 # gaussian diffusion trainer class
@@ -366,7 +368,9 @@ class GaussianDiffusion(nn.Module):
         timesteps = 1000,
         loss_type = 'l1',
         objective = 'pred_noise',
-        beta_schedule = 'cosine'
+        beta_schedule = 'cosine',
+        p2_loss_weight_gamma = 0., # p2 loss weight, from https://arxiv.org/abs/2204.00227 - 0 is equivalent to weight of 1 across time - 1. is recommended
+        p2_loss_weight_k = 1
     ):
         super().__init__()
         assert not (type(self) == GaussianDiffusion and denoise_fn.channels != denoise_fn.out_dim)
@@ -421,6 +425,10 @@ class GaussianDiffusion(nn.Module):
         register_buffer('posterior_mean_coef1', betas * torch.sqrt(alphas_cumprod_prev) / (1. - alphas_cumprod))
         register_buffer('posterior_mean_coef2', (1. - alphas_cumprod_prev) * torch.sqrt(alphas) / (1. - alphas_cumprod))
 
+        # calculate p2 reweighting
+
+        register_buffer('p2_loss_weight', (p2_loss_weight_k + alphas_cumprod / (1 - alphas_cumprod)) ** -p2_loss_weight_gamma)
+
     def predict_start_from_noise(self, x_t, t, noise):
         return (
             extract(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t -
@@ -527,8 +535,11 @@ class GaussianDiffusion(nn.Module):
         else:
             raise ValueError(f'unknown objective {self.objective}')
 
-        loss = self.loss_fn(model_out, target)
-        return loss
+        loss = self.loss_fn(model_out, target, reduction = 'none')
+        loss = reduce(loss, 'b ... -> b (...)', 'mean')
+
+        loss = loss * extract(self.p2_loss_weight, t, loss.shape)
+        return loss.mean()
 
     def forward(self, img, *args, **kwargs):
         b, c, h, w, device, img_size, = *img.shape, img.device, self.image_size
@@ -541,7 +552,7 @@ class GaussianDiffusion(nn.Module):
 # dataset classes
 
 class Dataset(data.Dataset):
-    def __init__(self, folder, image_size, exts = ['jpg', 'jpeg', 'png']):
+    def __init__(self, folder, image_size, exts = ['jpg', 'jpeg', 'png'], augment_horizontal_flip = False):
         super().__init__()
         self.folder = folder
         self.image_size = image_size
@@ -549,7 +560,7 @@ class Dataset(data.Dataset):
 
         self.transform = transforms.Compose([
             transforms.Resize(image_size),
-            transforms.RandomHorizontalFlip(),
+            transforms.RandomHorizontalFlip() if augment_horizontal_flip else nn.Identity(),
             transforms.CenterCrop(image_size),
             transforms.ToTensor()
         ])
@@ -580,7 +591,8 @@ class Trainer(object):
         step_start_ema = 2000,
         update_ema_every = 10,
         save_and_sample_every = 1000,
-        results_folder = './results'
+        results_folder = './results',
+        augment_horizontal_flip = True
     ):
         super().__init__()
         self.model = diffusion_model
@@ -596,8 +608,8 @@ class Trainer(object):
         self.gradient_accumulate_every = gradient_accumulate_every
         self.train_num_steps = train_num_steps
 
-        self.ds = Dataset(folder, image_size)
-        self.dl = cycle(data.DataLoader(self.ds, batch_size = train_batch_size, shuffle=True, pin_memory=True))
+        self.ds = Dataset(folder, image_size, augment_horizontal_flip = augment_horizontal_flip)
+        self.dl = cycle(data.DataLoader(self.ds, batch_size = train_batch_size, shuffle = True, pin_memory = True, num_workers = cpu_count()))
         self.opt = Adam(diffusion_model.parameters(), lr=train_lr)
 
         self.step = 0

setup.py

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