image is in range of 0 to 1

README.md

@@ -2,9 +2,7 @@
 
 ## Denoising Diffusion Probabilistic Model, in Pytorch
 
-Implementation of <a href="https://arxiv.org/abs/2006.11239">Denoising Diffusion Probabilistic Model</a> in Pytorch. It is a new approach to generative modeling that may <a href="https://ajolicoeur.wordpress.com/the-new-contender-to-gans-score-matching-with-langevin-sampling/">have the potential</a> to rival GANs. It uses denoising score matching to estimate the gradient of the data distribution, followed by Langevin sampling to sample from the true distribution.
+Implementation of <a href="https://arxiv.org/abs/2006.11239">Denoising Diffusion Probabilistic Model</a> in Pytorch. It is a new approach to generative modeling that may <a href="https://ajolicoeur.wordpress.com/the-new-contender-to-gans-score-matching-with-langevin-sampling/">have the potential</a> to rival GANs. It uses denoising score matching to estimate the gradient of the data distribution, followed by Langevin sampling to sample from the true distribution. This implementation was transcribed from the official Tensorflow version <a href="https://github.com/hojonathanho/diffusion">here</a>. 
-
-This implementation was transcribed from the official Tensorflow version <a href="https://github.com/hojonathanho/diffusion">here</a>
 
 <img src="./sample.png" width="500px"><img>
 
@@ -34,7 +32,7 @@ diffusion = GaussianDiffusion(
     loss_type = 'l1'    # L1 or L2
 )
 
-training_images = torch.randn(8, 3, 128, 128) # images are normalized from 0 to 1
+training_images = torch.randn(8, 3, 128, 128)
 loss = diffusion(training_images)
 loss.backward()
 # after a lot of training
@@ -64,11 +62,11 @@ trainer = Trainer(
     diffusion,
     'path/to/your/images',
     train_batch_size = 32,
-    train_lr = 1e-4,
+    train_lr = 2e-5,
     train_num_steps = 700000,         # total training steps
     gradient_accumulate_every = 2,    # gradient accumulation steps
     ema_decay = 0.995,                # exponential moving average decay
-    amp = True                        # turn on mixed precision
+    fp16 = True                       # turn on mixed precision training with apex
 )
 
 trainer.train()
@@ -79,43 +77,23 @@ Samples and model checkpoints will be logged to `./results` periodically
 ## Citations
 
 ```bibtex
-@inproceedings{NEURIPS2020_4c5bcfec,
+@misc{ho2020denoising,
-    author      = {Ho, Jonathan and Jain, Ajay and Abbeel, Pieter},
+    title   = {Denoising Diffusion Probabilistic Models},
-    booktitle   = {Advances in Neural Information Processing Systems},
+    author  = {Jonathan Ho and Ajay Jain and Pieter Abbeel},
-    editor      = {H. Larochelle and M. Ranzato and R. Hadsell and M.F. Balcan and H. Lin},
+    year    = {2020},
-    pages       = {6840--6851},
+    eprint  = {2006.11239},
-    publisher   = {Curran Associates, Inc.},
+    archivePrefix = {arXiv},
-    title       = {Denoising Diffusion Probabilistic Models},
+    primaryClass = {cs.LG}
-    url         = {https://proceedings.neurips.cc/paper/2020/file/4c5bcfec8584af0d967f1ab10179ca4b-Paper.pdf},
-    volume      = {33},
-    year        = {2020}
 }
 ```
 
 ```bibtex
-@InProceedings{pmlr-v139-nichol21a,
+@inproceedings{anonymous2021improved,
-    title       = {Improved Denoising Diffusion Probabilistic Models},
+    title   = {Improved Denoising Diffusion Probabilistic Models},
-    author      = {Nichol, Alexander Quinn and Dhariwal, Prafulla},
+    author  = {Anonymous},
-    booktitle   = {Proceedings of the 38th International Conference on Machine Learning},
+    booktitle = {Submitted to International Conference on Learning Representations},
-    pages       = {8162--8171},
+    year    = {2021},
-    year        = {2021},
+    url     = {https://openreview.net/forum?id=-NEXDKk8gZ},
-    editor      = {Meila, Marina and Zhang, Tong},
+    note    = {under review}
-    volume      = {139},
-    series      = {Proceedings of Machine Learning Research},
-    month       = {18--24 Jul},
-    publisher   = {PMLR},
-    pdf         = {http://proceedings.mlr.press/v139/nichol21a/nichol21a.pdf},
-    url         = {https://proceedings.mlr.press/v139/nichol21a.html},
-}
-```
-
-```bibtex
-@inproceedings{kingma2021on,
-    title       = {On Density Estimation with Diffusion Models},
-    author      = {Diederik P Kingma and Tim Salimans and Ben Poole and Jonathan Ho},
-    booktitle   = {Advances in Neural Information Processing Systems},
-    editor      = {A. Beygelzimer and Y. Dauphin and P. Liang and J. Wortman Vaughan},
-    year        = {2021},
-    url         = {https://openreview.net/forum?id=2LdBqxc1Yv}
 }
 ```

denoising_diffusion_pytorch/__init__.py

@@ -1,5 +1 @@
 from denoising_diffusion_pytorch.denoising_diffusion_pytorch import GaussianDiffusion, Unet, Trainer
-
-from denoising_diffusion_pytorch.learned_gaussian_diffusion import LearnedGaussianDiffusion
-from denoising_diffusion_pytorch.continuous_time_gaussian_diffusion import ContinuousTimeGaussianDiffusion
-from denoising_diffusion_pytorch.weighted_objective_gaussian_diffusion import WeightedObjectiveGaussianDiffusion

denoising_diffusion_pytorch/continuous_time_gaussian_diffusion.py

@@ -1,191 +0,0 @@
-import torch
-from torch import sqrt
-from torch import nn, einsum
-import torch.nn.functional as F
-from torch.special import expm1
-
-from tqdm import tqdm
-from einops import rearrange, repeat
-
-# helpers
-
-def exists(val):
-    return val is not None
-
-def default(val, d):
-    if exists(val):
-        return val
-    return d() if callable(d) else d
-
-# normalization functions
-
-def normalize_to_neg_one_to_one(img):
-    return img * 2 - 1
-
-def unnormalize_to_zero_to_one(t):
-    return (t + 1) * 0.5
-
-# diffusion helpers
-
-def right_pad_dims_to(x, t):
-    padding_dims = x.ndim - t.ndim
-    if padding_dims <= 0:
-        return t
-    return t.view(*t.shape, *((1,) * padding_dims))
-
-# continuous schedules
-
-# equations are taken from https://openreview.net/attachment?id=2LdBqxc1Yv&name=supplementary_material
-# @crowsonkb Katherine's repository also helped here https://github.com/crowsonkb/v-diffusion-jax/blob/master/diffusion/utils.py
-
-# log(snr) that approximates the original linear schedule
-
-def beta_linear_log_snr(t):
-    return -torch.log(expm1(1e-4 + 10 * (t ** 2)))
-
-def alpha_cosine_log_snr(t):
-    raise NotImplementedError
-
-class learned_noise_schedule(nn.Module):
-    def __init__(self):
-        super().__init__()
-        raise NotImplementedError
-        # learned noise schedule, using learned monotonic MLP (weights kept positive) in the paper
-
-class ContinuousTimeGaussianDiffusion(nn.Module):
-    def __init__(
-        self,
-        denoise_fn,
-        *,
-        image_size,
-        channels = 3,
-        cond_scale = 500,
-        loss_type = 'l1',
-        noise_schedule = 'linear',
-        num_sample_steps = 500
-    ):
-        super().__init__()
-        assert not denoise_fn.sinusoidal_cond_mlp
-
-        self.denoise_fn = denoise_fn
-
-        # image dimensions
-
-        self.channels = channels
-        self.image_size = image_size
-
-        # continuous noise schedule related stuff
-
-        self.cond_scale = cond_scale # the log(snr) will be scaled by this value
-        self.loss_type = loss_type
-
-        if noise_schedule == 'linear':
-            self.log_snr = beta_linear_log_snr
-        else:
-            raise ValueError(f'unknown noise schedule {noise_schedule}')
-
-        # sampling
-
-        self.num_sample_steps = num_sample_steps
-
-    @property
-    def device(self):
-        return next(self.denoise_fn.parameters()).device
-
-    @property
-    def loss_fn(self):
-        if self.loss_type == 'l1':
-            return F.l1_loss
-        elif self.loss_type == 'l2':
-            return F.mse_loss
-        else:
-            raise ValueError(f'invalid loss type {self.loss_type}')
-
-    def p_mean_variance(self, x, time, time_next):
-        # reviewer found an error in the equation in the paper (missing sigma)
-        # following - https://openreview.net/forum?id=2LdBqxc1Yv&noteId=rIQgH0zKsRt
-
-        # todo - derive x_start from the posterior mean and do dynamic thresholding
-        # assumed that is what is going on in Imagen
-
-        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
-
-        return model_mean, posterior_variance
-
-    # sampling related functions
-
-    @torch.no_grad()
-    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)
-
-        if time_next == 0:
-            return model_mean
-
-        noise = torch.randn_like(x)
-        return model_mean + sqrt(model_variance) * noise
-
-    @torch.no_grad()
-    def p_sample_loop(self, shape):
-        batch = shape[0]
-
-        img = torch.randn(shape, device = self.device)
-        steps = torch.linspace(1., 0., self.num_sample_steps + 1, device = self.device)
-
-        for i in tqdm(range(self.num_sample_steps), desc = 'sampling loop time step', total = self.num_sample_steps):
-            times = steps[i]
-            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
-
-    @torch.no_grad()
-    def sample(self, batch_size = 16):
-        return self.p_sample_loop((batch_size, self.channels, self.image_size, self.image_size))
-
-    # training related functions - noise prediction
-
-    def q_sample(self, x_start, times, noise = None):
-        noise = default(noise, lambda: torch.randn_like(x_start))
-
-        log_snr = self.log_snr(times)
-
-        log_snr_padded = right_pad_dims_to(x_start, log_snr)
-        alpha, sigma = sqrt(log_snr_padded.sigmoid()), sqrt((-log_snr_padded).sigmoid())
-        x_noised =  x_start * alpha + noise * sigma
-
-        return x_noised, log_snr
-
-    def random_times(self, batch_size):
-        # times are now uniform from 0 to 1
-        return torch.zeros((batch_size,), device = self.device).float().uniform_(0, 1)
-
-    def p_losses(self, x_start, times, noise = None):
-        noise = default(noise, lambda: torch.randn_like(x_start))
-
-        x, log_snr = self.q_sample(x_start = x_start, times = times, noise = noise)
-
-        model_out = self.denoise_fn(x, log_snr * self.cond_scale)
-        return self.loss_fn(model_out, noise)
-
-    def forward(self, img, *args, **kwargs):
-        b, c, h, w, device, img_size, = *img.shape, img.device, self.image_size
-        assert h == img_size and w == img_size, f'height and width of image must be {img_size}'
-
-        times = self.random_times(b)
-        img = normalize_to_neg_one_to_one(img)
-        return self.p_losses(img, times, *args, **kwargs)

denoising_diffusion_pytorch/denoising_diffusion_pytorch.py

@@ -7,16 +7,29 @@ from inspect import isfunction
 from functools import partial
 
 from torch.utils import data
-from torch.cuda.amp import autocast, GradScaler
-
 from pathlib import Path
 from torch.optim import Adam
 from torchvision import transforms, utils
 from PIL import Image
 
+import numpy as np
 from tqdm import tqdm
 from einops import rearrange
-from einops.layers.torch import Rearrange
+
+try:
+    from apex import amp
+    APEX_AVAILABLE = True
+except:
+    APEX_AVAILABLE = False
+
+# constants
+
+SAVE_AND_SAMPLE_EVERY = 1000
+UPDATE_EMA_EVERY = 10
+EXTS = ['jpg', 'jpeg', 'png']
+
+RESULTS_FOLDER = Path('./results')
+RESULTS_FOLDER.mkdir(exist_ok = True)
 
 # helpers functions
 
@@ -41,11 +54,12 @@ def num_to_groups(num, divisor):
         arr.append(remainder)
     return arr
 
-def normalize_to_neg_one_to_one(img):
+def loss_backwards(fp16, loss, optimizer, **kwargs):
-    return img * 2 - 1
+    if fp16:
-
+        with amp.scale_loss(loss, optimizer) as scaled_loss:
-def unnormalize_to_zero_to_one(t):
+            scaled_loss.backward(**kwargs)
-    return (t + 1) * 0.5
+    else:
+        loss.backward(**kwargs)
 
 # small helper modules
 
@@ -86,111 +100,69 @@ class SinusoidalPosEmb(nn.Module):
         emb = torch.cat((emb.sin(), emb.cos()), dim=-1)
         return emb
 
-def Upsample(dim):
+class Mish(nn.Module):
-    return nn.ConvTranspose2d(dim, dim, 4, 2, 1)
+    def forward(self, x):
+        return x * torch.tanh(F.softplus(x))
 
-def Downsample(dim):
+class Upsample(nn.Module):
-    return nn.Conv2d(dim, dim, 4, 2, 1)
+    def __init__(self, dim):
-
-class LayerNorm(nn.Module):
-    def __init__(self, dim, eps = 1e-5):
         super().__init__()
-        self.eps = eps
+        self.conv = nn.ConvTranspose2d(dim, dim, 4, 2, 1)
-        self.g = nn.Parameter(torch.ones(1, dim, 1, 1))
-        self.b = nn.Parameter(torch.zeros(1, dim, 1, 1))
 
     def forward(self, x):
-        var = torch.var(x, dim = 1, unbiased = False, keepdim = True)
+        return self.conv(x)
-        mean = torch.mean(x, dim = 1, keepdim = True)
-        return (x - mean) / (var + self.eps).sqrt() * self.g + self.b
 
-class PreNorm(nn.Module):
+class Downsample(nn.Module):
-    def __init__(self, dim, fn):
+    def __init__(self, dim):
+        super().__init__()
+        self.conv = nn.Conv2d(dim, dim, 3, 2, 1)
+
+    def forward(self, x):
+        return self.conv(x)
+
+class Rezero(nn.Module):
+    def __init__(self, fn):
         super().__init__()
         self.fn = fn
-        self.norm = LayerNorm(dim)
+        self.g = nn.Parameter(torch.zeros(1))
 
     def forward(self, x):
-        x = self.norm(x)
+        return self.fn(x) * self.g
-        return self.fn(x)
 
 # building block modules
 
 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.block = nn.Sequential(
-        self.norm = nn.GroupNorm(groups, dim_out)
+            nn.Conv2d(dim, dim_out, 3, padding=1),
-        self.act = nn.SiLU()
+            nn.GroupNorm(groups, dim_out),
-
+            Mish()
-    def forward(self, x, scale_shift = None):
+        )
-        x = self.proj(x)
+    def forward(self, x):
-        x = self.norm(x)
+        return self.block(x)
-
-        if exists(scale_shift):
-            scale, shift = scale_shift
-            x = x * (scale + 1) + shift
-
-        x = self.act(x)
-        return x
 
 class ResnetBlock(nn.Module):
-    def __init__(self, dim, dim_out, *, time_emb_dim = None, groups = 8):
+    def __init__(self, dim, dim_out, *, time_emb_dim, groups = 8):
         super().__init__()
         self.mlp = nn.Sequential(
-            nn.SiLU(),
+            Mish(),
-            nn.Linear(time_emb_dim, dim_out * 2)
+            nn.Linear(time_emb_dim, dim_out)
-        ) if exists(time_emb_dim) else None
+        )
 
-        self.block1 = Block(dim, dim_out, groups = groups)
+        self.block1 = Block(dim, dim_out)
-        self.block2 = Block(dim_out, dim_out, groups = groups)
+        self.block2 = Block(dim_out, dim_out)
         self.res_conv = nn.Conv2d(dim, dim_out, 1) if dim != dim_out else nn.Identity()
 
-    def forward(self, x, time_emb = None):
+    def forward(self, x, time_emb):
-
+        h = self.block1(x)
-        scale_shift = None
+        h += self.mlp(time_emb)[:, :, None, None]
-        if exists(self.mlp) and exists(time_emb):
-            time_emb = self.mlp(time_emb)
-            time_emb = rearrange(time_emb, 'b c -> b c 1 1')
-            scale_shift = time_emb.chunk(2, dim = 1)
-
-        h = self.block1(x, scale_shift = scale_shift)
-
         h = self.block2(h)
         return h + self.res_conv(x)
 
 class LinearAttention(nn.Module):
     def __init__(self, dim, heads = 4, dim_head = 32):
         super().__init__()
-        self.scale = dim_head ** -0.5
-        self.heads = heads
-        hidden_dim = dim_head * heads
-        self.to_qkv = nn.Conv2d(dim, hidden_dim * 3, 1, bias = False)
-
-        self.to_out = nn.Sequential(
-            nn.Conv2d(hidden_dim, dim, 1),
-            LayerNorm(dim)
-        )
-
-    def forward(self, x):
-        b, c, h, w = x.shape
-        qkv = self.to_qkv(x).chunk(3, dim = 1)
-        q, k, v = map(lambda t: rearrange(t, 'b (h c) x y -> b h c (x y)', h = self.heads), qkv)
-
-        q = q.softmax(dim = -2)
-        k = k.softmax(dim = -1)
-
-        q = q * self.scale
-        context = torch.einsum('b h d n, b h e n -> b h d e', k, v)
-
-        out = torch.einsum('b h d e, b h d n -> b h e n', context, q)
-        out = rearrange(out, 'b h c (x y) -> b (h c) x y', h = self.heads, x = h, y = w)
-        return self.to_out(out)
-
-class Attention(nn.Module):
-    def __init__(self, dim, heads = 4, dim_head = 32):
-        super().__init__()
-        self.scale = dim_head ** -0.5
         self.heads = heads
         hidden_dim = dim_head * heads
         self.to_qkv = nn.Conv2d(dim, hidden_dim * 3, 1, bias = False)
@@ -198,75 +170,37 @@ class Attention(nn.Module):
 
     def forward(self, x):
         b, c, h, w = x.shape
-        qkv = self.to_qkv(x).chunk(3, dim = 1)
+        qkv = self.to_qkv(x)
-        q, k, v = map(lambda t: rearrange(t, 'b (h c) x y -> b h c (x y)', h = self.heads), qkv)
+        q, k, v = rearrange(qkv, 'b (qkv heads c) h w -> qkv b heads c (h w)', heads = self.heads, qkv=3)
-        q = q * self.scale
+        k = k.softmax(dim=-1)
-
+        context = torch.einsum('bhdn,bhen->bhde', k, v)
-        sim = einsum('b h d i, b h d j -> b h i j', q, k)
+        out = torch.einsum('bhde,bhdn->bhen', context, q)
-        sim = sim - sim.amax(dim = -1, keepdim = True).detach()
+        out = rearrange(out, 'b heads c (h w) -> b (heads c) h w', heads=self.heads, h=h, w=w)
-        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)
 
 # model
 
-def MLP(dim_in, dim_hidden):
-    return nn.Sequential(
-        Rearrange('... -> ... 1'),
-        nn.Linear(1, dim_hidden),
-        nn.GELU(),
-        nn.LayerNorm(dim_hidden),
-        nn.Linear(dim_hidden, dim_hidden),
-        nn.GELU(),
-        nn.LayerNorm(dim_hidden),
-        nn.Linear(dim_hidden, dim_hidden)
-    )
-
 class Unet(nn.Module):
     def __init__(
         self,
         dim,
-        init_dim = None,
         out_dim = None,
         dim_mults=(1, 2, 4, 8),
-        channels = 3,
+        groups = 8,
-        resnet_block_groups = 8,
+        channels = 3
-        learned_variance = False,
-        sinusoidal_cond_mlp = True
     ):
         super().__init__()
-
-        # determine dimensions
-
         self.channels = channels
 
-        init_dim = default(init_dim, dim // 3 * 2)
+        dims = [channels, *map(lambda m: dim * m, dim_mults)]
-        self.init_conv = nn.Conv2d(channels, init_dim, 7, padding = 3)
-
-        dims = [init_dim, *map(lambda m: dim * m, dim_mults)]
         in_out = list(zip(dims[:-1], dims[1:]))
 
-        block_klass = partial(ResnetBlock, groups = resnet_block_groups)
+        self.time_pos_emb = SinusoidalPosEmb(dim)
-
+        self.mlp = nn.Sequential(
-        # time embeddings
+            nn.Linear(dim, dim * 4),
-
+            Mish(),
-        time_dim = dim * 4
+            nn.Linear(dim * 4, dim)
-
+        )
-        self.sinusoidal_cond_mlp = sinusoidal_cond_mlp
-
-        if sinusoidal_cond_mlp:
-            self.time_mlp = nn.Sequential(
-                SinusoidalPosEmb(dim),
-                nn.Linear(dim, time_dim),
-                nn.GELU(),
-                nn.Linear(time_dim, time_dim)
-            )
-        else:
-            self.time_mlp = MLP(1, time_dim)
-
-        # layers
 
         self.downs = nn.ModuleList([])
         self.ups = nn.ModuleList([])
@@ -276,44 +210,42 @@ class Unet(nn.Module):
             is_last = ind >= (num_resolutions - 1)
 
             self.downs.append(nn.ModuleList([
-                block_klass(dim_in, dim_out, time_emb_dim = time_dim),
+                ResnetBlock(dim_in, dim_out, time_emb_dim = dim),
-                block_klass(dim_out, dim_out, time_emb_dim = time_dim),
+                ResnetBlock(dim_out, dim_out, time_emb_dim = dim),
-                Residual(PreNorm(dim_out, LinearAttention(dim_out))),
+                Residual(Rezero(LinearAttention(dim_out))),
                 Downsample(dim_out) if not is_last else nn.Identity()
             ]))
 
         mid_dim = dims[-1]
-        self.mid_block1 = block_klass(mid_dim, mid_dim, time_emb_dim = time_dim)
+        self.mid_block1 = ResnetBlock(mid_dim, mid_dim, time_emb_dim = dim)
-        self.mid_attn = Residual(PreNorm(mid_dim, Attention(mid_dim)))
+        self.mid_attn = Residual(Rezero(LinearAttention(mid_dim)))
-        self.mid_block2 = block_klass(mid_dim, mid_dim, time_emb_dim = time_dim)
+        self.mid_block2 = ResnetBlock(mid_dim, mid_dim, time_emb_dim = dim)
 
         for ind, (dim_in, dim_out) in enumerate(reversed(in_out[1:])):
             is_last = ind >= (num_resolutions - 1)
 
             self.ups.append(nn.ModuleList([
-                block_klass(dim_out * 2, dim_in, time_emb_dim = time_dim),
+                ResnetBlock(dim_out * 2, dim_in, time_emb_dim = dim),
-                block_klass(dim_in, dim_in, time_emb_dim = time_dim),
+                ResnetBlock(dim_in, dim_in, time_emb_dim = dim),
-                Residual(PreNorm(dim_in, LinearAttention(dim_in))),
+                Residual(Rezero(LinearAttention(dim_in))),
                 Upsample(dim_in) if not is_last else nn.Identity()
             ]))
 
-        default_out_dim = channels * (1 if not learned_variance else 2)
+        out_dim = default(out_dim, channels)
-        self.out_dim = default(out_dim, default_out_dim)
-
         self.final_conv = nn.Sequential(
-            block_klass(dim, dim),
+            Block(dim, dim),
-            nn.Conv2d(dim, self.out_dim, 1)
+            nn.Conv2d(dim, out_dim, 1)
         )
 
     def forward(self, x, time):
-        x = self.init_conv(x)
+        t = self.time_pos_emb(time)
-        t = self.time_mlp(time)
+        t = self.mlp(t)
 
         h = []
 
-        for block1, block2, attn, downsample in self.downs:
+        for resnet, resnet2, attn, downsample in self.downs:
-            x = block1(x, t)
+            x = resnet(x, t)
-            x = block2(x, t)
+            x = resnet2(x, t)
             x = attn(x)
             h.append(x)
             x = downsample(x)
@@ -322,10 +254,10 @@ class Unet(nn.Module):
         x = self.mid_attn(x)
         x = self.mid_block2(x, t)
 
-        for block1, block2, attn, upsample in self.ups:
+        for resnet, resnet2, attn, upsample in self.ups:
             x = torch.cat((x, h.pop()), dim=1)
-            x = block1(x, t)
+            x = resnet(x, t)
-            x = block2(x, t)
+            x = resnet2(x, t)
             x = attn(x)
             x = upsample(x)
 
@@ -338,11 +270,10 @@ def extract(a, t, x_shape):
     out = a.gather(-1, t)
     return out.reshape(b, *((1,) * (len(x_shape) - 1)))
 
-def linear_beta_schedule(timesteps):
+def noise_like(shape, device, repeat=False):
-    scale = 1000 / timesteps
+    repeat_noise = lambda: torch.randn((1, *shape[1:]), device=device).repeat(shape[0], *((1,) * (len(shape) - 1)))
-    beta_start = scale * 0.0001
+    noise = lambda: torch.randn(shape, device=device)
-    beta_end = scale * 0.02
+    return repeat_noise() if repeat else noise()
-    return torch.linspace(beta_start, beta_end, timesteps, dtype = torch.float64)
 
 def cosine_beta_schedule(timesteps, s = 0.008):
     """
@@ -350,11 +281,11 @@ def cosine_beta_schedule(timesteps, s = 0.008):
     as proposed in https://openreview.net/forum?id=-NEXDKk8gZ
     """
     steps = timesteps + 1
-    x = torch.linspace(0, timesteps, steps, dtype = torch.float64)
+    x = np.linspace(0, steps, steps)
-    alphas_cumprod = torch.cos(((x / timesteps) + s) / (1 + s) * torch.pi * 0.5) ** 2
+    alphas_cumprod = np.cos(((x / steps) + s) / (1 + s) * np.pi * 0.5) ** 2
     alphas_cumprod = alphas_cumprod / alphas_cumprod[0]
     betas = 1 - (alphas_cumprod[1:] / alphas_cumprod[:-1])
-    return torch.clip(betas, 0, 0.999)
+    return np.clip(betas, a_min = 0, a_max = 0.999)
 
 class GaussianDiffusion(nn.Module):
     def __init__(
@@ -365,61 +296,55 @@ class GaussianDiffusion(nn.Module):
         channels = 3,
         timesteps = 1000,
         loss_type = 'l1',
-        objective = 'pred_noise',
+        betas = None
-        beta_schedule = 'cosine'
     ):
         super().__init__()
-        assert not (type(self) == GaussianDiffusion and denoise_fn.channels != denoise_fn.out_dim)
-
         self.channels = channels
         self.image_size = image_size
         self.denoise_fn = denoise_fn
-        self.objective = objective
 
-        if beta_schedule == 'linear':
+        if exists(betas):
-            betas = linear_beta_schedule(timesteps)
+            betas = betas.detach().cpu().numpy() if isinstance(betas, torch.Tensor) else betas
-        elif beta_schedule == 'cosine':
-            betas = cosine_beta_schedule(timesteps)
         else:
-            raise ValueError(f'unknown beta schedule {beta_schedule}')
+            betas = cosine_beta_schedule(timesteps)
 
         alphas = 1. - betas
-        alphas_cumprod = torch.cumprod(alphas, axis=0)
+        alphas_cumprod = np.cumprod(alphas, axis=0)
-        alphas_cumprod_prev = F.pad(alphas_cumprod[:-1], (1, 0), value = 1.)
+        alphas_cumprod_prev = np.append(1., alphas_cumprod[:-1])
 
         timesteps, = betas.shape
         self.num_timesteps = int(timesteps)
         self.loss_type = loss_type
 
-        # helper function to register buffer from float64 to float32
+        to_torch = partial(torch.tensor, dtype=torch.float32)
 
-        register_buffer = lambda name, val: self.register_buffer(name, val.to(torch.float32))
+        self.register_buffer('betas', to_torch(betas))
-
+        self.register_buffer('alphas_cumprod', to_torch(alphas_cumprod))
-        register_buffer('betas', betas)
+        self.register_buffer('alphas_cumprod_prev', to_torch(alphas_cumprod_prev))
-        register_buffer('alphas_cumprod', alphas_cumprod)
-        register_buffer('alphas_cumprod_prev', alphas_cumprod_prev)
 
         # calculations for diffusion q(x_t | x_{t-1}) and others
-
+        self.register_buffer('sqrt_alphas_cumprod', to_torch(np.sqrt(alphas_cumprod)))
-        register_buffer('sqrt_alphas_cumprod', torch.sqrt(alphas_cumprod))
+        self.register_buffer('sqrt_one_minus_alphas_cumprod', to_torch(np.sqrt(1. - alphas_cumprod)))
-        register_buffer('sqrt_one_minus_alphas_cumprod', torch.sqrt(1. - alphas_cumprod))
+        self.register_buffer('log_one_minus_alphas_cumprod', to_torch(np.log(1. - alphas_cumprod)))
-        register_buffer('log_one_minus_alphas_cumprod', torch.log(1. - alphas_cumprod))
+        self.register_buffer('sqrt_recip_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod)))
-        register_buffer('sqrt_recip_alphas_cumprod', torch.sqrt(1. / alphas_cumprod))
+        self.register_buffer('sqrt_recipm1_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod - 1)))
-        register_buffer('sqrt_recipm1_alphas_cumprod', torch.sqrt(1. / alphas_cumprod - 1))
 
         # calculations for posterior q(x_{t-1} | x_t, x_0)
-
         posterior_variance = betas * (1. - alphas_cumprod_prev) / (1. - alphas_cumprod)
-
         # above: equal to 1. / (1. / (1. - alpha_cumprod_tm1) + alpha_t / beta_t)
-
+        self.register_buffer('posterior_variance', to_torch(posterior_variance))
-        register_buffer('posterior_variance', posterior_variance)
-
         # below: log calculation clipped because the posterior variance is 0 at the beginning of the diffusion chain
+        self.register_buffer('posterior_log_variance_clipped', to_torch(np.log(np.maximum(posterior_variance, 1e-20))))
+        self.register_buffer('posterior_mean_coef1', to_torch(
+            betas * np.sqrt(alphas_cumprod_prev) / (1. - alphas_cumprod)))
+        self.register_buffer('posterior_mean_coef2', to_torch(
+            (1. - alphas_cumprod_prev) * np.sqrt(alphas) / (1. - alphas_cumprod)))
 
-        register_buffer('posterior_log_variance_clipped', torch.log(posterior_variance.clamp(min =1e-20)))
+    def q_mean_variance(self, x_start, t):
-        register_buffer('posterior_mean_coef1', betas * torch.sqrt(alphas_cumprod_prev) / (1. - alphas_cumprod))
+        mean = extract(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start
-        register_buffer('posterior_mean_coef2', (1. - alphas_cumprod_prev) * torch.sqrt(alphas) / (1. - alphas_cumprod))
+        variance = extract(1. - self.alphas_cumprod, t, x_start.shape)
+        log_variance = extract(self.log_one_minus_alphas_cumprod, t, x_start.shape)
+        return mean, variance, log_variance
 
     def predict_start_from_noise(self, x_t, t, noise):
         return (
@@ -437,26 +362,19 @@ class GaussianDiffusion(nn.Module):
         return posterior_mean, posterior_variance, posterior_log_variance_clipped
 
     def p_mean_variance(self, x, t, clip_denoised: bool):
-        model_output = self.denoise_fn(x, t)
+        x_recon = self.predict_start_from_noise(x, t=t, noise=self.denoise_fn(x, t))
-
-        if self.objective == 'pred_noise':
-            x_start = self.predict_start_from_noise(x, t = t, noise = model_output)
-        elif self.objective == 'pred_x0':
-            x_start = model_output
-        else:
-            raise ValueError(f'unknown objective {self.objective}')
 
         if clip_denoised:
-            x_start.clamp_(-1., 1.)
+            x_recon.clamp_(0., 1.)
 
-        model_mean, posterior_variance, posterior_log_variance = self.q_posterior(x_start = x_start, x_t = x, t = t)
+        model_mean, posterior_variance, posterior_log_variance = self.q_posterior(x_start=x_recon, x_t=x, t=t)
         return model_mean, posterior_variance, posterior_log_variance
 
     @torch.no_grad()
-    def p_sample(self, x, t, clip_denoised=True):
+    def p_sample(self, x, t, clip_denoised=True, repeat_noise=False):
         b, *_, device = *x.shape, x.device
         model_mean, _, model_log_variance = self.p_mean_variance(x=x, t=t, clip_denoised=clip_denoised)
-        noise = torch.randn_like(x)
+        noise = noise_like(x.shape, device, repeat_noise)
         # no noise when t == 0
         nonzero_mask = (1 - (t == 0).float()).reshape(b, *((1,) * (len(x.shape) - 1)))
         return model_mean + nonzero_mask * (0.5 * model_log_variance).exp() * noise
@@ -470,8 +388,6 @@ class GaussianDiffusion(nn.Module):
 
         for i in tqdm(reversed(range(0, self.num_timesteps)), desc='sampling loop time step', total=self.num_timesteps):
             img = self.p_sample(img, torch.full((b,), i, device=device, dtype=torch.long))
-
-        img = unnormalize_to_zero_to_one(img)
         return img
 
     @torch.no_grad()
@@ -504,48 +420,36 @@ class GaussianDiffusion(nn.Module):
             extract(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape) * noise
         )
 
-    @property
-    def loss_fn(self):
-        if self.loss_type == 'l1':
-            return F.l1_loss
-        elif self.loss_type == 'l2':
-            return F.mse_loss
-        else:
-            raise ValueError(f'invalid loss type {self.loss_type}')
-
     def p_losses(self, x_start, t, noise = None):
         b, c, h, w = x_start.shape
         noise = default(noise, lambda: torch.randn_like(x_start))
 
-        x = self.q_sample(x_start=x_start, t=t, noise=noise)
+        x_noisy = self.q_sample(x_start=x_start, t=t, noise=noise)
-        model_out = self.denoise_fn(x, t)
+        x_recon = self.denoise_fn(x_noisy, t)
 
-        if self.objective == 'pred_noise':
+        if self.loss_type == 'l1':
-            target = noise
+            loss = (noise - x_recon).abs().mean()
-        elif self.objective == 'pred_x0':
+        elif self.loss_type == 'l2':
-            target = x_start
+            loss = F.mse_loss(noise, x_recon)
         else:
-            raise ValueError(f'unknown objective {self.objective}')
+            raise NotImplementedError()
 
-        loss = self.loss_fn(model_out, target)
         return loss
 
-    def forward(self, img, *args, **kwargs):
+    def forward(self, x, *args, **kwargs):
-        b, c, h, w, device, img_size, = *img.shape, img.device, self.image_size
+        b, c, h, w, device, img_size, = *x.shape, x.device, self.image_size
         assert h == img_size and w == img_size, f'height and width of image must be {img_size}'
         t = torch.randint(0, self.num_timesteps, (b,), device=device).long()
-
+        return self.p_losses(x, t, *args, **kwargs)
-        img = normalize_to_neg_one_to_one(img)
-        return self.p_losses(img, t, *args, **kwargs)
 
 # dataset classes
 
 class Dataset(data.Dataset):
-    def __init__(self, folder, image_size, exts = ['jpg', 'jpeg', 'png']):
+    def __init__(self, folder, image_size):
         super().__init__()
         self.folder = folder
         self.image_size = image_size
-        self.paths = [p for ext in exts for p in Path(f'{folder}').glob(f'**/*.{ext}')]
+        self.paths = [p for ext in EXTS for p in Path(f'{folder}').glob(f'**/*.{ext}')]
 
         self.transform = transforms.Compose([
             transforms.Resize(image_size),
@@ -573,23 +477,17 @@ class Trainer(object):
         ema_decay = 0.995,
         image_size = 128,
         train_batch_size = 32,
-        train_lr = 1e-4,
+        train_lr = 2e-5,
         train_num_steps = 100000,
         gradient_accumulate_every = 2,
-        amp = False,
+        fp16 = False,
-        step_start_ema = 2000,
+        step_start_ema = 2000
-        update_ema_every = 10,
-        save_and_sample_every = 1000,
-        results_folder = './results'
     ):
         super().__init__()
         self.model = diffusion_model
         self.ema = EMA(ema_decay)
         self.ema_model = copy.deepcopy(self.model)
-        self.update_ema_every = update_ema_every
-
         self.step_start_ema = step_start_ema
-        self.save_and_sample_every = save_and_sample_every
 
         self.batch_size = train_batch_size
         self.image_size = diffusion_model.image_size
@@ -602,11 +500,11 @@ class Trainer(object):
 
         self.step = 0
 
-        self.amp = amp
+        assert not fp16 or fp16 and APEX_AVAILABLE, 'Apex must be installed in order for mixed precision training to be turned on'
-        self.scaler = GradScaler(enabled = amp)
 
-        self.results_folder = Path(results_folder)
+        self.fp16 = fp16
-        self.results_folder.mkdir(exist_ok = True)
+        if fp16:
+            (self.model, self.ema_model), self.opt = amp.initialize([self.model, self.ema_model], self.opt, opt_level='O1')
 
         self.reset_parameters()
 
@@ -623,50 +521,41 @@ class Trainer(object):
         data = {
             'step': self.step,
             'model': self.model.state_dict(),
-            'ema': self.ema_model.state_dict(),
+            'ema': self.ema_model.state_dict()
-            'scaler': self.scaler.state_dict()
         }
-        torch.save(data, str(self.results_folder / f'model-{milestone}.pt'))
+        torch.save(data, str(RESULTS_FOLDER / f'model-{milestone}.pt'))
 
     def load(self, milestone):
-        data = torch.load(str(self.results_folder / f'model-{milestone}.pt'))
+        data = torch.load(str(RESULTS_FOLDER / f'model-{milestone}.pt'))
 
         self.step = data['step']
         self.model.load_state_dict(data['model'])
         self.ema_model.load_state_dict(data['ema'])
-        self.scaler.load_state_dict(data['scaler'])
 
     def train(self):
-        with tqdm(initial = self.step, total = self.train_num_steps) as pbar:
+        backwards = partial(loss_backwards, self.fp16)
 
-            while self.step < self.train_num_steps:
+        while self.step < self.train_num_steps:
-                for i in range(self.gradient_accumulate_every):
+            for i in range(self.gradient_accumulate_every):
-                    data = next(self.dl).cuda()
+                data = next(self.dl).cuda()
+                loss = self.model(data)
+                print(f'{self.step}: {loss.item()}')
+                backwards(loss / self.gradient_accumulate_every, self.opt)
 
-                    with autocast(enabled = self.amp):
+            self.opt.step()
-                        loss = self.model(data)
+            self.opt.zero_grad()
-                        self.scaler.scale(loss / self.gradient_accumulate_every).backward()
 
-                    pbar.set_description(f'loss: {loss.item():.4f}')
+            if self.step % UPDATE_EMA_EVERY == 0:
+                self.step_ema()
 
-                self.scaler.step(self.opt)
+            if self.step != 0 and self.step % SAVE_AND_SAMPLE_EVERY == 0:
-                self.scaler.update()
+                milestone = self.step // SAVE_AND_SAMPLE_EVERY
-                self.opt.zero_grad()
+                batches = num_to_groups(36, self.batch_size)
+                all_images_list = list(map(lambda n: self.ema_model.sample(batch_size=n), batches))
+                all_images = torch.cat(all_images_list, dim=0)
+                utils.save_image(all_images, str(RESULTS_FOLDER / f'sample-{milestone}.png'), nrow=6)
+                self.save(milestone)
 
-                if self.step % self.update_ema_every == 0:
+            self.step += 1
-                    self.step_ema()
 
-                if self.step != 0 and self.step % self.save_and_sample_every == 0:
+        print('training completed')
-                    self.ema_model.eval()
-
-                    milestone = self.step // self.save_and_sample_every
-                    batches = num_to_groups(36, self.batch_size)
-                    all_images_list = list(map(lambda n: self.ema_model.sample(batch_size=n), batches))
-                    all_images = torch.cat(all_images_list, dim=0)
-                    utils.save_image(all_images, str(self.results_folder / f'sample-{milestone}.png'), nrow = 6)
-                    self.save(milestone)
-
-                self.step += 1
-                pbar.update(1)
-
-        print('training complete')

denoising_diffusion_pytorch/learned_gaussian_diffusion.py

@@ -1,132 +0,0 @@
-import torch
-from math import pi, sqrt, log as ln
-from inspect import isfunction
-from torch import nn, einsum
-from einops import rearrange
-
-from denoising_diffusion_pytorch.denoising_diffusion_pytorch import GaussianDiffusion, extract, unnormalize_to_zero_to_one
-
-# constants
-
-NAT = 1. / ln(2)
-
-# helper functions
-
-def exists(x):
-    return x is not None
-
-def default(val, d):
-    if exists(val):
-        return val
-    return d() if isfunction(d) else d
-
-# tensor helpers
-
-def log(t, eps = 1e-12):
-    return torch.log(t.clamp(min = eps))
-
-def meanflat(x):
-    return x.mean(dim = tuple(range(1, len(x.shape))))
-
-def normal_kl(mean1, logvar1, mean2, logvar2):
-    """
-    KL divergence between normal distributions parameterized by mean and log-variance.
-    """
-    return 0.5 * (-1.0 + logvar2 - logvar1 + torch.exp(logvar1 - logvar2) + ((mean1 - mean2) ** 2) * torch.exp(-logvar2))
-
-def approx_standard_normal_cdf(x):
-    return 0.5 * (1.0 + torch.tanh(sqrt(2.0 / pi) * (x + 0.044715 * (x ** 3))))
-
-def discretized_gaussian_log_likelihood(x, *, means, log_scales, thres = 0.999):
-    assert x.shape == means.shape == log_scales.shape
-
-    centered_x = x - means
-    inv_stdv = torch.exp(-log_scales)
-    plus_in = inv_stdv * (centered_x + 1. / 255.)
-    cdf_plus = approx_standard_normal_cdf(plus_in)
-    min_in = inv_stdv * (centered_x - 1. / 255.)
-    cdf_min = approx_standard_normal_cdf(min_in)
-    log_cdf_plus = log(cdf_plus)
-    log_one_minus_cdf_min = log(1. - cdf_min)
-    cdf_delta = cdf_plus - cdf_min
-
-    log_probs = torch.where(x < -thres,
-        log_cdf_plus,
-        torch.where(x > thres,
-            log_one_minus_cdf_min,
-            log(cdf_delta)))
-
-    return log_probs
-
-# https://arxiv.org/abs/2102.09672
-
-# i thought the results were questionable, if one were to focus only on FID
-# but may as well get this in here for others to try, as GLIDE is using it (and DALL-E2 first stage of cascade)
-# gaussian diffusion for learned variance + hybrid eps simple + vb loss
-
-class LearnedGaussianDiffusion(GaussianDiffusion):
-    def __init__(
-        self,
-        denoise_fn,
-        vb_loss_weight = 0.001,  # lambda was 0.001 in the paper
-        *args,
-        **kwargs
-    ):
-        super().__init__(denoise_fn, *args, **kwargs)
-        assert denoise_fn.out_dim == (denoise_fn.channels * 2), 'dimension out of unet must be twice the number of channels for learned variance - you can also set the `learned_variance` keyword argument on the Unet to be `True`'
-        self.vb_loss_weight = vb_loss_weight
-
-    def p_mean_variance(self, *, x, t, clip_denoised, model_output = None):
-        model_output = default(model_output, lambda: self.denoise_fn(x, t))
-        pred_noise, var_interp_frac_unnormalized = model_output.chunk(2, dim = 1)
-
-        min_log = extract(self.posterior_log_variance_clipped, t, x.shape)
-        max_log = extract(torch.log(self.betas), t, x.shape)
-        var_interp_frac = unnormalize_to_zero_to_one(var_interp_frac_unnormalized)
-
-        model_log_variance = var_interp_frac * max_log + (1 - var_interp_frac) * min_log
-        model_variance = model_log_variance.exp()
-
-        x_start = self.predict_start_from_noise(x, t, pred_noise)
-
-        if clip_denoised:
-            x_start.clamp_(-1., 1.)
-
-        model_mean, _, _ = self.q_posterior(x_start, x, t)
-
-        return model_mean, model_variance, model_log_variance
-
-    def p_losses(self, x_start, t, noise = None, clip_denoised = False):
-        noise = default(noise, lambda: torch.randn_like(x_start))
-        x_t = self.q_sample(x_start = x_start, t = t, noise = noise)
-
-        # model output
-
-        model_output = self.denoise_fn(x_t, t)
-
-        # calculating kl loss for learned variance (interpolation)
-
-        true_mean, _, true_log_variance_clipped = self.q_posterior(x_start = x_start, x_t = x_t, t = t)
-        model_mean, _, model_log_variance = self.p_mean_variance(x = x_t, t = t, clip_denoised = clip_denoised, model_output = model_output)
-
-        # kl loss with detached model predicted mean, for stability reasons as in paper
-
-        detached_model_mean = model_mean.detach()
-
-        kl = normal_kl(true_mean, true_log_variance_clipped, detached_model_mean, model_log_variance)
-        kl = meanflat(kl) * NAT
-
-        decoder_nll = -discretized_gaussian_log_likelihood(x_start, means = detached_model_mean, log_scales = 0.5 * model_log_variance)
-        decoder_nll = meanflat(decoder_nll) * NAT
-
-        # at the first timestep return the decoder NLL, otherwise return KL(q(x_{t-1}|x_t,x_0) || p(x_{t-1}|x_t))
-
-        vb_losses = torch.where(t == 0, decoder_nll, kl)
-
-        # simple loss - predicting noise, x0, or x_prev
-
-        pred_noise, _ = model_output.chunk(2, dim = 1)
-
-        simple_losses = self.loss_fn(pred_noise, noise)
-
-        return simple_losses + vb_losses.mean() * self.vb_loss_weight

denoising_diffusion_pytorch/weighted_objective_gaussian_diffusion.py

@@ -1,80 +0,0 @@
-import torch
-from inspect import isfunction
-from torch import nn, einsum
-from einops import rearrange
-
-from denoising_diffusion_pytorch.denoising_diffusion_pytorch import GaussianDiffusion
-
-# helper functions
-
-def exists(x):
-    return x is not None
-
-def default(val, d):
-    if exists(val):
-        return val
-    return d() if isfunction(d) else d
-
-# some improvisation on my end
-# where i have the model learn to both predict noise and x0
-# and learn the weighted sum for each depending on time step
-
-class WeightedObjectiveGaussianDiffusion(GaussianDiffusion):
-    def __init__(
-        self,
-        denoise_fn,
-        *args,
-        pred_noise_loss_weight = 0.1,
-        pred_x_start_loss_weight = 0.1,
-        **kwargs
-    ):
-        super().__init__(denoise_fn, *args, **kwargs)
-        channels = denoise_fn.channels
-        assert denoise_fn.out_dim == (channels * 2 + 2), 'dimension out (out_dim) of unet must be twice the number of channels + 2 (for the softmax weighted sum) - for channels of 3, this should be (3 * 2) + 2 = 8'
-
-        self.split_dims = (channels, channels, 2)
-        self.pred_noise_loss_weight = pred_noise_loss_weight
-        self.pred_x_start_loss_weight = pred_x_start_loss_weight
-
-    def p_mean_variance(self, *, x, t, clip_denoised, model_output = None):
-        model_output = self.denoise_fn(x, t)
-
-        pred_noise, pred_x_start, weights = model_output.split(self.split_dims, dim = 1)
-        normalized_weights = weights.softmax(dim = 1)
-
-        x_start_from_noise = self.predict_start_from_noise(x, t = t, noise = pred_noise)
-        
-        x_starts = torch.stack((x_start_from_noise, pred_x_start), dim = 1)
-        weighted_x_start = einsum('b j h w, b j c h w -> b c h w', normalized_weights, x_starts)
-
-        if clip_denoised:
-            weighted_x_start.clamp_(-1., 1.)
-
-        model_mean, model_variance, model_log_variance = self.q_posterior(weighted_x_start, x, t)
-
-        return model_mean, model_variance, model_log_variance
-
-    def p_losses(self, x_start, t, noise = None, clip_denoised = False):
-        noise = default(noise, lambda: torch.randn_like(x_start))
-        x_t = self.q_sample(x_start = x_start, t = t, noise = noise)
-
-        model_output = self.denoise_fn(x_t, t)
-        pred_noise, pred_x_start, weights = model_output.split(self.split_dims, dim = 1)
-
-        # get loss for predicted noise and x_start
-        # with the loss weight given at initialization
-
-        noise_loss = self.loss_fn(noise, pred_noise) * self.pred_noise_loss_weight
-        x_start_loss = self.loss_fn(x_start, pred_x_start) * self.pred_x_start_loss_weight
-
-        # calculate x_start from predicted noise
-        # then do a weighted sum of the x_start prediction, weights also predicted by the model (softmax normalized)
-
-        x_start_from_pred_noise = self.predict_start_from_noise(x_t, t, pred_noise)
-        x_start_from_pred_noise = x_start_from_pred_noise.clamp(-2., 2.)
-        weighted_x_start = einsum('b j h w, b j c h w -> b c h w', weights.softmax(dim = 1), torch.stack((x_start_from_pred_noise, pred_x_start), dim = 1))
-
-        # main loss to x_start with the weighted one
-
-        weighted_x_start_loss = self.loss_fn(x_start, weighted_x_start)
-        return weighted_x_start_loss + x_start_loss + noise_loss

setup.py

@@ -3,19 +3,19 @@ from setuptools import setup, find_packages
 setup(
   name = 'denoising-diffusion-pytorch',
   packages = find_packages(),
-  version = '0.16.6',
+  version = '0.6.2',
   license='MIT',
   description = 'Denoising Diffusion Probabilistic Models - Pytorch',
   author = 'Phil Wang',
   author_email = 'lucidrains@gmail.com',
   url = 'https://github.com/lucidrains/denoising-diffusion-pytorch',
-  long_description_content_type = 'text/markdown',
   keywords = [
     'artificial intelligence',
     'generative models'
   ],
   install_requires=[
     'einops',
+    'numpy',
     'pillow',
     'torch',
     'torchvision',