bring in ddim sampling

Solve issue #26

README.md

@@ -2,7 +2,13 @@
 
 ## 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. This implementation was transcribed from the official Tensorflow version <a href="https://github.com/hojonathanho/diffusion">here</a>. 
+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>
+
+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>
 
@@ -32,7 +38,7 @@ diffusion = GaussianDiffusion(
     loss_type = 'l1'    # L1 or L2
 )
 
-training_images = torch.randn(8, 3, 128, 128)
+training_images = torch.randn(8, 3, 128, 128) # images are normalized from 0 to 1
 loss = diffusion(training_images)
 loss.backward()
 # after a lot of training
@@ -54,19 +60,20 @@ model = Unet(
 diffusion = GaussianDiffusion(
     model,
     image_size = 128,
-    timesteps = 1000,   # number of steps
-    loss_type = 'l1'    # L1 or L2
+    timesteps = 1000,           # number of steps
+    sampling_timesteps = 250,   # number of sampling timesteps (using ddim for faster inference [see citation for ddim paper])
+    loss_type = 'l1'            # L1 or L2
 ).cuda()
 
 trainer = Trainer(
     diffusion,
     'path/to/your/images',
     train_batch_size = 32,
-    train_lr = 2e-5,
+    train_lr = 1e-4,
     train_num_steps = 700000,         # total training steps
     gradient_accumulate_every = 2,    # gradient accumulation steps
     ema_decay = 0.995,                # exponential moving average decay
-    fp16 = True                       # turn on mixed precision training with apex
+    amp = True                        # turn on mixed precision
 )
 
 trainer.train()
@@ -74,26 +81,92 @@ trainer.train()
 
 Samples and model checkpoints will be logged to `./results` periodically
 
+## Multi-GPU Training
+
+The `Trainer` class is now equipped with <a href="https://huggingface.co/docs/accelerate/accelerator">🤗 Accelerator</a>. You can easily do multi-gpu training in two steps using their `accelerate` CLI
+
+At the project root directory, where the training script is, run
+
+```python
+$ accelerate config
+```
+
+Then, in the same directory
+
+```python
+$ accelerate launch train.py
+```
+
 ## Citations
 
 ```bibtex
-@misc{ho2020denoising,
-    title   = {Denoising Diffusion Probabilistic Models},
-    author  = {Jonathan Ho and Ajay Jain and Pieter Abbeel},
-    year    = {2020},
-    eprint  = {2006.11239},
-    archivePrefix = {arXiv},
-    primaryClass = {cs.LG}
+@inproceedings{NEURIPS2020_4c5bcfec,
+    author      = {Ho, Jonathan and Jain, Ajay and Abbeel, Pieter},
+    booktitle   = {Advances in Neural Information Processing Systems},
+    editor      = {H. Larochelle and M. Ranzato and R. Hadsell and M.F. Balcan and H. Lin},
+    pages       = {6840--6851},
+    publisher   = {Curran Associates, Inc.},
+    title       = {Denoising Diffusion Probabilistic Models},
+    url         = {https://proceedings.neurips.cc/paper/2020/file/4c5bcfec8584af0d967f1ab10179ca4b-Paper.pdf},
+    volume      = {33},
+    year        = {2020}
 }
 ```
 
 ```bibtex
-@inproceedings{anonymous2021improved,
-    title   = {Improved Denoising Diffusion Probabilistic Models},
-    author  = {Anonymous},
-    booktitle = {Submitted to International Conference on Learning Representations},
+@InProceedings{pmlr-v139-nichol21a,
+    title       = {Improved Denoising Diffusion Probabilistic Models},
+    author      = {Nichol, Alexander Quinn and Dhariwal, Prafulla},
+    booktitle   = {Proceedings of the 38th International Conference on Machine Learning},
+    pages       = {8162--8171},
+    year        = {2021},
+    editor      = {Meila, Marina and Zhang, Tong},
+    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}
+}
+```
+
+```bibtex
+@article{Choi2022PerceptionPT,
+    title   = {Perception Prioritized Training of Diffusion Models},
+    author  = {Jooyoung Choi and Jungbeom Lee and Chaehun Shin and Sungwon Kim and Hyunwoo J. Kim and Sung-Hoon Yoon},
+    journal = {ArXiv},
+    year    = {2022},
+    volume  = {abs/2204.00227}
+}
+```
+
+```bibtex
+@article{Karras2022ElucidatingTD,
+    title   = {Elucidating the Design Space of Diffusion-Based Generative Models},
+    author  = {Tero Karras and Miika Aittala and Timo Aila and Samuli Laine},
+    journal = {ArXiv},
+    year    = {2022},
+    volume  = {abs/2206.00364}
+}
+```
+
+```bibtex
+@article{Song2021DenoisingDI,
+    title   = {Denoising Diffusion Implicit Models},
+    author  = {Jiaming Song and Chenlin Meng and Stefano Ermon},
+    journal = {ArXiv},
     year    = {2021},
-    url     = {https://openreview.net/forum?id=-NEXDKk8gZ},
-    note    = {under review}
+    volume  = {abs/2010.02502}
 }
 ```

denoising_diffusion_pytorch/__init__.py

@@ -1 +1,6 @@
 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
+from denoising_diffusion_pytorch.elucidated_diffusion import ElucidatedDiffusion

denoising_diffusion_pytorch/continuous_time_gaussian_diffusion.py

@@ -0,0 +1,287 @@
+import math
+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, reduce
+from einops.layers.torch import Rearrange
+
+# 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))
+
+# neural net helpers
+
+class Residual(nn.Module):
+    def __init__(self, fn):
+        super().__init__()
+        self.fn = fn
+
+    def forward(self, x):
+        return x + self.fn(x)
+
+class MonotonicLinear(nn.Module):
+    def __init__(self, *args, **kwargs):
+        super().__init__()
+        self.net = nn.Linear(*args, **kwargs)
+
+    def forward(self, x):
+        return F.linear(x, self.net.weight.abs(), self.net.bias.abs())
+
+# 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 log(t, eps = 1e-20):
+    return torch.log(t.clamp(min = eps))
+
+def beta_linear_log_snr(t):
+    return -log(expm1(1e-4 + 10 * (t ** 2)))
+
+def alpha_cosine_log_snr(t, s = 0.008):
+    return -log((torch.cos((t + s) / (1 + s) * math.pi * 0.5) ** -2) - 1, eps = 1e-5)
+
+class learned_noise_schedule(nn.Module):
+    """ described in section H and then I.2 of the supplementary material for variational ddpm paper """
+
+    def __init__(
+        self,
+        *,
+        log_snr_max,
+        log_snr_min,
+        hidden_dim = 1024,
+        frac_gradient = 1.
+    ):
+        super().__init__()
+        self.slope = log_snr_min - log_snr_max
+        self.intercept = log_snr_max
+
+        self.net = nn.Sequential(
+            Rearrange('... -> ... 1'),
+            MonotonicLinear(1, 1),
+            Residual(nn.Sequential(
+                MonotonicLinear(1, hidden_dim),
+                nn.Sigmoid(),
+                MonotonicLinear(hidden_dim, 1)
+            )),
+            Rearrange('... 1 -> ...'),
+        )
+
+        self.frac_gradient = frac_gradient
+
+    def forward(self, x):
+        frac_gradient = self.frac_gradient
+        device = x.device
+
+        out_zero = self.net(torch.zeros_like(x))
+        out_one =  self.net(torch.ones_like(x))
+
+        x = self.net(x)
+
+        normed = self.slope * ((x - out_zero) / (out_one - out_zero)) + self.intercept
+        return normed * frac_gradient + normed.detach() * (1 - frac_gradient)
+
+class ContinuousTimeGaussianDiffusion(nn.Module):
+    def __init__(
+        self,
+        model,
+        *,
+        image_size,
+        channels = 3,
+        loss_type = 'l1',
+        noise_schedule = 'linear',
+        num_sample_steps = 500,
+        clip_sample_denoised = True,
+        learned_schedule_net_hidden_dim = 1024,
+        learned_noise_schedule_frac_gradient = 1.,   # between 0 and 1, determines what percentage of gradients go back, so one can update the learned noise schedule more slowly
+        p2_loss_weight_gamma = 0.,                   # p2 loss weight, from https://arxiv.org/abs/2204.00227 - 0 is equivalent to weight of 1 across time
+        p2_loss_weight_k = 1
+    ):
+        super().__init__()
+        assert model.learned_sinusoidal_cond
+
+        self.model = model
+
+        # image dimensions
+
+        self.channels = channels
+        self.image_size = image_size
+
+        # continuous noise schedule related stuff
+
+        self.loss_type = loss_type
+
+        if noise_schedule == 'linear':
+            self.log_snr = beta_linear_log_snr
+        elif noise_schedule == 'cosine':
+            self.log_snr = alpha_cosine_log_snr
+        elif noise_schedule == 'learned':
+            log_snr_max, log_snr_min = [beta_linear_log_snr(torch.tensor([time])).item() for time in (0., 1.)]
+
+            self.log_snr = learned_noise_schedule(
+                log_snr_max = log_snr_max,
+                log_snr_min = log_snr_min,
+                hidden_dim = learned_schedule_net_hidden_dim,
+                frac_gradient = learned_noise_schedule_frac_gradient
+            )
+        else:
+            raise ValueError(f'unknown noise schedule {noise_schedule}')
+
+        # sampling
+
+        self.num_sample_steps = num_sample_steps
+        self.clip_sample_denoised = clip_sample_denoised
+
+        # p2 loss weight
+        # proposed https://arxiv.org/abs/2204.00227
+
+        assert p2_loss_weight_gamma <= 2, 'in paper, they noticed any gamma greater than 2 is harmful'
+
+        self.p2_loss_weight_gamma = p2_loss_weight_gamma  # recommended to be 0.5 or 1
+        self.p2_loss_weight_k = p2_loss_weight_k
+
+    @property
+    def device(self):
+        return next(self.model.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
+
+        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()
+
+        alpha, sigma, alpha_next = map(sqrt, (squared_alpha, squared_sigma, squared_alpha_next))
+
+        batch_log_snr = repeat(log_snr, ' -> b', b = x.shape[0])
+        pred_noise = self.model(x, batch_log_snr)
+
+        if self.clip_sample_denoised:
+            x_start = (x - sigma * pred_noise) / alpha
+
+            # in Imagen, this was changed to dynamic thresholding
+            x_start.clamp_(-1., 1.)
+
+            model_mean = alpha_next * (x * (1 - c) / alpha + c * x_start)
+        else:
+            model_mean = alpha_next / alpha * (x - c * 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.model(x, log_snr)
+
+        losses = self.loss_fn(model_out, noise, reduction = 'none')
+        losses = reduce(losses, 'b ... -> b', 'mean')
+
+        if self.p2_loss_weight_gamma >= 0:
+            # following eq 8. in https://arxiv.org/abs/2204.00227
+            loss_weight = (self.p2_loss_weight_k + log_snr.exp()) ** -self.p2_loss_weight_gamma
+            losses = losses * loss_weight
+
+        return losses.mean()
+
+    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

@@ -4,32 +4,28 @@ import torch
 from torch import nn, einsum
 import torch.nn.functional as F
 from inspect import isfunction
+from collections import namedtuple
 from functools import partial
 
-from torch.utils import data
+from torch.utils.data import Dataset, DataLoader
+from multiprocessing import cpu_count
+
 from pathlib import Path
 from torch.optim import Adam
-from torchvision import transforms, utils
+from torchvision import transforms as T, utils
 from PIL import Image
 
-import numpy as np
-from tqdm import tqdm
-from einops import rearrange
+from einops import rearrange, reduce
+from einops.layers.torch import Rearrange
 
-try:
-    from apex import amp
-    APEX_AVAILABLE = True
-except:
-    APEX_AVAILABLE = False
+from tqdm.auto import tqdm
+from ema_pytorch import EMA
+
+from accelerate import Accelerator
 
 # constants
 
-SAVE_AND_SAMPLE_EVERY = 1000
-UPDATE_EMA_EVERY = 10
-EXTS = ['jpg', 'jpeg', 'png']
-
-RESULTS_FOLDER = Path('./results')
-RESULTS_FOLDER.mkdir(exist_ok = True)
+ModelPrediction =  namedtuple('ModelPrediction', ['pred_noise', 'pred_x_start'])
 
 # helpers functions
 
@@ -46,6 +42,9 @@ def cycle(dl):
         for data in dl:
             yield data
 
+def has_int_squareroot(num):
+    return (math.sqrt(num) ** 2) == num
+
 def num_to_groups(num, divisor):
     groups = num // divisor
     remainder = num % divisor
@@ -54,30 +53,21 @@ def num_to_groups(num, divisor):
         arr.append(remainder)
     return arr
 
-def loss_backwards(fp16, loss, optimizer, **kwargs):
-    if fp16:
-        with amp.scale_loss(loss, optimizer) as scaled_loss:
-            scaled_loss.backward(**kwargs)
-    else:
-        loss.backward(**kwargs)
+def convert_image_to(img_type, image):
+    if image.mode != img_type:
+        return image.convert(img_type)
+    return image
+
+# 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
 
 # small helper modules
 
-class EMA():
-    def __init__(self, beta):
-        super().__init__()
-        self.beta = beta
-
-    def update_model_average(self, ma_model, current_model):
-        for current_params, ma_params in zip(current_model.parameters(), ma_model.parameters()):
-            old_weight, up_weight = ma_params.data, current_params.data
-            ma_params.data = self.update_average(old_weight, up_weight)
-
-    def update_average(self, old, new):
-        if old is None:
-            return new
-        return old * self.beta + (1 - self.beta) * new
-
 class Residual(nn.Module):
     def __init__(self, fn):
         super().__init__()
@@ -86,6 +76,39 @@ class Residual(nn.Module):
     def forward(self, x, *args, **kwargs):
         return self.fn(x, *args, **kwargs) + x
 
+def Upsample(dim, dim_out = None):
+    return nn.Sequential(
+        nn.Upsample(scale_factor = 2, mode = 'nearest'),
+        nn.Conv2d(dim, default(dim_out, dim), 3, padding = 1)
+    )
+
+def Downsample(dim, dim_out = None):
+    return nn.Conv2d(dim, default(dim_out, dim), 4, 2, 1)
+
+class LayerNorm(nn.Module):
+    def __init__(self, dim, eps = 1e-5):
+        super().__init__()
+        self.eps = eps
+        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)
+        mean = torch.mean(x, dim = 1, keepdim = True)
+        return (x - mean) / (var + self.eps).sqrt() * self.g + self.b
+
+class PreNorm(nn.Module):
+    def __init__(self, dim, fn):
+        super().__init__()
+        self.fn = fn
+        self.norm = LayerNorm(dim)
+
+    def forward(self, x):
+        x = self.norm(x)
+        return self.fn(x)
+
+# sinusoidal positional embeds
+
 class SinusoidalPosEmb(nn.Module):
     def __init__(self, dim):
         super().__init__()
@@ -100,69 +123,101 @@ class SinusoidalPosEmb(nn.Module):
         emb = torch.cat((emb.sin(), emb.cos()), dim=-1)
         return emb
 
-class Mish(nn.Module):
-    def forward(self, x):
-        return x * torch.tanh(F.softplus(x))
+class LearnedSinusoidalPosEmb(nn.Module):
+    """ following @crowsonkb 's lead with learned sinusoidal pos emb """
+    """ https://github.com/crowsonkb/v-diffusion-jax/blob/master/diffusion/models/danbooru_128.py#L8 """
 
-class Upsample(nn.Module):
     def __init__(self, dim):
         super().__init__()
-        self.conv = nn.ConvTranspose2d(dim, dim, 4, 2, 1)
+        assert (dim % 2) == 0
+        half_dim = dim // 2
+        self.weights = nn.Parameter(torch.randn(half_dim))
 
     def forward(self, x):
-        return self.conv(x)
-
-class Downsample(nn.Module):
-    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.g = nn.Parameter(torch.zeros(1))
-
-    def forward(self, x):
-        return self.fn(x) * self.g
+        x = rearrange(x, 'b -> b 1')
+        freqs = x * rearrange(self.weights, 'd -> 1 d') * 2 * math.pi
+        fouriered = torch.cat((freqs.sin(), freqs.cos()), dim = -1)
+        fouriered = torch.cat((x, fouriered), dim = -1)
+        return fouriered
 
 # building block modules
 
 class Block(nn.Module):
     def __init__(self, dim, dim_out, groups = 8):
         super().__init__()
-        self.block = nn.Sequential(
-            nn.Conv2d(dim, dim_out, 3, padding=1),
-            nn.GroupNorm(groups, dim_out),
-            Mish()
-        )
-    def forward(self, x):
-        return self.block(x)
+        self.proj = nn.Conv2d(dim, dim_out, 3, padding = 1)
+        self.norm = nn.GroupNorm(groups, dim_out)
+        self.act = nn.SiLU()
+
+    def forward(self, x, scale_shift = None):
+        x = self.proj(x)
+        x = self.norm(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, groups = 8):
+    def __init__(self, dim, dim_out, *, time_emb_dim = None, groups = 8):
         super().__init__()
         self.mlp = nn.Sequential(
-            Mish(),
-            nn.Linear(time_emb_dim, dim_out)
-        )
+            nn.SiLU(),
+            nn.Linear(time_emb_dim, dim_out * 2)
+        ) if exists(time_emb_dim) else None
 
-        self.block1 = Block(dim, dim_out)
-        self.block2 = Block(dim_out, dim_out)
+        self.block1 = Block(dim, dim_out, groups = groups)
+        self.block2 = Block(dim_out, dim_out, groups = groups)
         self.res_conv = nn.Conv2d(dim, dim_out, 1) if dim != dim_out else nn.Identity()
 
-    def forward(self, x, time_emb):
-        h = self.block1(x)
-        h += self.mlp(time_emb)[:, :, None, None]
+    def forward(self, x, time_emb = None):
+
+        scale_shift = 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)
@@ -170,12 +225,16 @@ class LinearAttention(nn.Module):
 
     def forward(self, x):
         b, c, h, w = x.shape
-        qkv = self.to_qkv(x)
-        q, k, v = rearrange(qkv, 'b (qkv heads c) h w -> qkv b heads c (h w)', heads = self.heads, qkv=3)
-        k = k.softmax(dim=-1)
-        context = torch.einsum('bhdn,bhen->bhde', k, v)
-        out = torch.einsum('bhde,bhdn->bhen', context, q)
-        out = rearrange(out, 'b heads c (h w) -> b (heads c) h w', heads=self.heads, h=h, w=w)
+        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 * self.scale
+
+        sim = einsum('b h d i, b h d j -> b h i j', q, k)
+        sim = sim - sim.amax(dim = -1, keepdim = True).detach()
+        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
@@ -184,24 +243,51 @@ class Unet(nn.Module):
     def __init__(
         self,
         dim,
+        init_dim = None,
         out_dim = None,
         dim_mults=(1, 2, 4, 8),
-        groups = 8,
-        channels = 3
+        channels = 3,
+        resnet_block_groups = 8,
+        learned_variance = False,
+        learned_sinusoidal_cond = False,
+        learned_sinusoidal_dim = 16
     ):
         super().__init__()
+
+        # determine dimensions
+
         self.channels = channels
 
-        dims = [channels, *map(lambda m: dim * m, dim_mults)]
+        init_dim = default(init_dim, dim)
+        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:]))
 
-        self.time_pos_emb = SinusoidalPosEmb(dim)
-        self.mlp = nn.Sequential(
-            nn.Linear(dim, dim * 4),
-            Mish(),
-            nn.Linear(dim * 4, dim)
+        block_klass = partial(ResnetBlock, groups = resnet_block_groups)
+
+        # time embeddings
+
+        time_dim = dim * 4
+
+        self.learned_sinusoidal_cond = learned_sinusoidal_cond
+
+        if learned_sinusoidal_cond:
+            sinu_pos_emb = LearnedSinusoidalPosEmb(learned_sinusoidal_dim)
+            fourier_dim = learned_sinusoidal_dim + 1
+        else:
+            sinu_pos_emb = SinusoidalPosEmb(dim)
+            fourier_dim = dim
+
+        self.time_mlp = nn.Sequential(
+            sinu_pos_emb,
+            nn.Linear(fourier_dim, time_dim),
+            nn.GELU(),
+            nn.Linear(time_dim, time_dim)
         )
 
+        # layers
+
         self.downs = nn.ModuleList([])
         self.ups = nn.ModuleList([])
         num_resolutions = len(in_out)
@@ -210,57 +296,68 @@ class Unet(nn.Module):
             is_last = ind >= (num_resolutions - 1)
 
             self.downs.append(nn.ModuleList([
-                ResnetBlock(dim_in, dim_out, time_emb_dim = dim),
-                ResnetBlock(dim_out, dim_out, time_emb_dim = dim),
-                Residual(Rezero(LinearAttention(dim_out))),
-                Downsample(dim_out) if not is_last else nn.Identity()
+                block_klass(dim_in, dim_in, time_emb_dim = time_dim),
+                block_klass(dim_in, dim_in, time_emb_dim = time_dim),
+                Residual(PreNorm(dim_in, LinearAttention(dim_in))),
+                Downsample(dim_in, dim_out) if not is_last else nn.Conv2d(dim_in, dim_out, 3, padding = 1)
             ]))
 
         mid_dim = dims[-1]
-        self.mid_block1 = ResnetBlock(mid_dim, mid_dim, time_emb_dim = dim)
-        self.mid_attn = Residual(Rezero(LinearAttention(mid_dim)))
-        self.mid_block2 = ResnetBlock(mid_dim, mid_dim, time_emb_dim = dim)
+        self.mid_block1 = block_klass(mid_dim, mid_dim, time_emb_dim = time_dim)
+        self.mid_attn = Residual(PreNorm(mid_dim, Attention(mid_dim)))
+        self.mid_block2 = block_klass(mid_dim, mid_dim, time_emb_dim = time_dim)
 
-        for ind, (dim_in, dim_out) in enumerate(reversed(in_out[1:])):
-            is_last = ind >= (num_resolutions - 1)
+        for ind, (dim_in, dim_out) in enumerate(reversed(in_out)):
+            is_last = ind == (len(in_out) - 1)
 
             self.ups.append(nn.ModuleList([
-                ResnetBlock(dim_out * 2, dim_in, time_emb_dim = dim),
-                ResnetBlock(dim_in, dim_in, time_emb_dim = dim),
-                Residual(Rezero(LinearAttention(dim_in))),
-                Upsample(dim_in) if not is_last else nn.Identity()
+                block_klass(dim_out + dim_in, dim_out, time_emb_dim = time_dim),
+                block_klass(dim_out + dim_in, dim_out, time_emb_dim = time_dim),
+                Residual(PreNorm(dim_out, LinearAttention(dim_out))),
+                Upsample(dim_out, dim_in) if not is_last else  nn.Conv2d(dim_out, dim_in, 3, padding = 1)
             ]))
 
-        out_dim = default(out_dim, channels)
-        self.final_conv = nn.Sequential(
-            Block(dim, dim),
-            nn.Conv2d(dim, out_dim, 1)
-        )
+        default_out_dim = channels * (1 if not learned_variance else 2)
+        self.out_dim = default(out_dim, default_out_dim)
+
+        self.final_res_block = block_klass(dim * 2, dim, time_emb_dim = time_dim)
+        self.final_conv = nn.Conv2d(dim, self.out_dim, 1)
 
     def forward(self, x, time):
-        t = self.time_pos_emb(time)
-        t = self.mlp(t)
+        x = self.init_conv(x)
+        r = x.clone()
+
+        t = self.time_mlp(time)
 
         h = []
 
-        for resnet, resnet2, attn, downsample in self.downs:
-            x = resnet(x, t)
-            x = resnet2(x, t)
+        for block1, block2, attn, downsample in self.downs:
+            x = block1(x, t)
+            h.append(x)
+
+            x = block2(x, t)
             x = attn(x)
             h.append(x)
+
             x = downsample(x)
 
         x = self.mid_block1(x, t)
         x = self.mid_attn(x)
         x = self.mid_block2(x, t)
 
-        for resnet, resnet2, attn, upsample in self.ups:
-            x = torch.cat((x, h.pop()), dim=1)
-            x = resnet(x, t)
-            x = resnet2(x, t)
+        for block1, block2, attn, upsample in self.ups:
+            x = torch.cat((x, h.pop()), dim = 1)
+            x = block1(x, t)
+
+            x = torch.cat((x, h.pop()), dim = 1)
+            x = block2(x, t)
             x = attn(x)
+
             x = upsample(x)
 
+        x = torch.cat((x, r), dim = 1)
+
+        x = self.final_res_block(x, t)
         return self.final_conv(x)
 
 # gaussian diffusion trainer class
@@ -270,10 +367,11 @@ def extract(a, t, x_shape):
     out = a.gather(-1, t)
     return out.reshape(b, *((1,) * (len(x_shape) - 1)))
 
-def noise_like(shape, device, repeat=False):
-    repeat_noise = lambda: torch.randn((1, *shape[1:]), device=device).repeat(shape[0], *((1,) * (len(shape) - 1)))
-    noise = lambda: torch.randn(shape, device=device)
-    return repeat_noise() if repeat else noise()
+def linear_beta_schedule(timesteps):
+    scale = 1000 / timesteps
+    beta_start = scale * 0.0001
+    beta_end = scale * 0.02
+    return torch.linspace(beta_start, beta_end, timesteps, dtype = torch.float64)
 
 def cosine_beta_schedule(timesteps, s = 0.008):
     """
@@ -281,70 +379,94 @@ def cosine_beta_schedule(timesteps, s = 0.008):
     as proposed in https://openreview.net/forum?id=-NEXDKk8gZ
     """
     steps = timesteps + 1
-    x = np.linspace(0, steps, steps)
-    alphas_cumprod = np.cos(((x / steps) + s) / (1 + s) * np.pi * 0.5) ** 2
+    x = torch.linspace(0, timesteps, steps, dtype = torch.float64)
+    alphas_cumprod = torch.cos(((x / timesteps) + s) / (1 + s) * math.pi * 0.5) ** 2
     alphas_cumprod = alphas_cumprod / alphas_cumprod[0]
     betas = 1 - (alphas_cumprod[1:] / alphas_cumprod[:-1])
-    return np.clip(betas, a_min = 0, a_max = 0.999)
+    return torch.clip(betas, 0, 0.999)
 
 class GaussianDiffusion(nn.Module):
     def __init__(
         self,
-        denoise_fn,
+        model,
         *,
         image_size,
         channels = 3,
         timesteps = 1000,
+        sampling_timesteps = None,
         loss_type = 'l1',
-        betas = None
+        objective = 'pred_noise',
+        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,
+        ddim_sampling_eta = 1.
     ):
         super().__init__()
+        assert not (type(self) == GaussianDiffusion and model.channels != model.out_dim)
+
         self.channels = channels
         self.image_size = image_size
-        self.denoise_fn = denoise_fn
+        self.model = model
+        self.objective = objective
 
-        if exists(betas):
-            betas = betas.detach().cpu().numpy() if isinstance(betas, torch.Tensor) else betas
-        else:
+        assert objective in {'pred_noise', 'pred_x0'}, 'objective must be either pred_noise (predict noise) or pred_x0 (predict image start)'
+
+        if beta_schedule == 'linear':
+            betas = linear_beta_schedule(timesteps)
+        elif beta_schedule == 'cosine':
             betas = cosine_beta_schedule(timesteps)
+        else:
+            raise ValueError(f'unknown beta schedule {beta_schedule}')
 
         alphas = 1. - betas
-        alphas_cumprod = np.cumprod(alphas, axis=0)
-        alphas_cumprod_prev = np.append(1., alphas_cumprod[:-1])
+        alphas_cumprod = torch.cumprod(alphas, axis=0)
+        alphas_cumprod_prev = F.pad(alphas_cumprod[:-1], (1, 0), value = 1.)
 
         timesteps, = betas.shape
         self.num_timesteps = int(timesteps)
         self.loss_type = loss_type
 
-        to_torch = partial(torch.tensor, dtype=torch.float32)
+        # sampling related parameters
 
-        self.register_buffer('betas', to_torch(betas))
-        self.register_buffer('alphas_cumprod', to_torch(alphas_cumprod))
-        self.register_buffer('alphas_cumprod_prev', to_torch(alphas_cumprod_prev))
+        self.sampling_timesteps = default(sampling_timesteps, timesteps) # default num sampling timesteps to number of timesteps at training
+
+        assert self.sampling_timesteps <= timesteps
+        self.is_ddim_sampling = self.sampling_timesteps < timesteps
+        self.ddim_sampling_eta = ddim_sampling_eta
+
+        # helper function to register buffer from float64 to float32
+
+        register_buffer = lambda name, val: self.register_buffer(name, val.to(torch.float32))
+
+        register_buffer('betas', betas)
+        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)))
-        self.register_buffer('sqrt_one_minus_alphas_cumprod', to_torch(np.sqrt(1. - alphas_cumprod)))
-        self.register_buffer('log_one_minus_alphas_cumprod', to_torch(np.log(1. - alphas_cumprod)))
-        self.register_buffer('sqrt_recip_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod)))
-        self.register_buffer('sqrt_recipm1_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod - 1)))
+
+        register_buffer('sqrt_alphas_cumprod', torch.sqrt(alphas_cumprod))
+        register_buffer('sqrt_one_minus_alphas_cumprod', torch.sqrt(1. - alphas_cumprod))
+        register_buffer('log_one_minus_alphas_cumprod', torch.log(1. - alphas_cumprod))
+        register_buffer('sqrt_recip_alphas_cumprod', torch.sqrt(1. / alphas_cumprod))
+        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))
-        # 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)))
 
-    def q_mean_variance(self, x_start, t):
-        mean = extract(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start
-        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
+        posterior_variance = betas * (1. - alphas_cumprod_prev) / (1. - alphas_cumprod)
+
+        # above: equal to 1. / (1. / (1. - alpha_cumprod_tm1) + alpha_t / beta_t)
+
+        register_buffer('posterior_variance', posterior_variance)
+
+        # below: log calculation clipped because the posterior variance is 0 at the beginning of the diffusion chain
+
+        register_buffer('posterior_log_variance_clipped', torch.log(posterior_variance.clamp(min =1e-20)))
+        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 (
@@ -352,6 +474,12 @@ class GaussianDiffusion(nn.Module):
             extract(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape) * noise
         )
 
+    def predict_noise_from_start(self, x_t, t, x0):
+        return (
+            (x0 - extract(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t) / \
+            extract(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape)
+        )
+
     def q_posterior(self, x_start, x_t, t):
         posterior_mean = (
             extract(self.posterior_mean_coef1, t, x_t.shape) * x_start +
@@ -361,40 +489,83 @@ class GaussianDiffusion(nn.Module):
         posterior_log_variance_clipped = extract(self.posterior_log_variance_clipped, t, x_t.shape)
         return posterior_mean, posterior_variance, posterior_log_variance_clipped
 
+    def model_predictions(self, x, t):
+        model_output = self.model(x, t)
+
+        if self.objective == 'pred_noise':
+            pred_noise = model_output
+            x_start = self.predict_start_from_noise(x, t, model_output)
+
+        elif self.objective == 'pred_x0':
+            pred_noise = self.predict_noise_from_start(x, t, model_output)
+            x_start = model_output
+
+        return ModelPrediction(pred_noise, x_start)
+
     def p_mean_variance(self, x, t, clip_denoised: bool):
-        x_recon = self.predict_start_from_noise(x, t=t, noise=self.denoise_fn(x, t))
+        preds = self.model_predictions(x, t)
+        x_start = preds.pred_x_start
 
         if clip_denoised:
-            x_recon.clamp_(-1., 1.)
+            x_start.clamp_(-1., 1.)
 
-        model_mean, posterior_variance, posterior_log_variance = self.q_posterior(x_start=x_recon, x_t=x, t=t)
+        model_mean, posterior_variance, posterior_log_variance = self.q_posterior(x_start = x_start, 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, repeat_noise=False):
+    def p_sample(self, x, t: int, clip_denoised = True):
         b, *_, device = *x.shape, x.device
-        model_mean, _, model_log_variance = self.p_mean_variance(x=x, t=t, clip_denoised=clip_denoised)
-        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
+        batched_times = torch.full((x.shape[0],), t, device = x.device, dtype = torch.long)
+        model_mean, _, model_log_variance = self.p_mean_variance(x = x, t = batched_times, clip_denoised = clip_denoised)
+        noise = torch.randn_like(x) if t > 0 else 0. # no noise if t == 0
+        return model_mean + (0.5 * model_log_variance).exp() * noise
 
     @torch.no_grad()
     def p_sample_loop(self, shape):
-        device = self.betas.device
+        batch, device = shape[0], self.betas.device
 
-        b = shape[0]
         img = torch.randn(shape, device=device)
 
-        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))
+        for t in tqdm(reversed(range(0, self.num_timesteps)), desc = 'sampling loop time step'):
+            img = self.p_sample(img, t)
+
+        img = unnormalize_to_zero_to_one(img)
+        return img
+
+    @torch.no_grad()
+    def ddim_sample(self, shape, clip_denoised = False):
+        batch, device, total_timesteps, sampling_timesteps, eta, objective = shape[0], self.betas.device, self.num_timesteps, self.sampling_timesteps, self.ddim_sampling_eta, self.objective
+
+        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:]))
+
+        img = torch.randn(shape, device = device)
+
+        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]
+
+            time_cond = torch.full((batch,), time, device = device, dtype = torch.long)
+
+            pred_noise, x_start, *_ = self.model_predictions(img, time_cond)
+
+            c1 = eta * ((1 - alpha / alpha_next) * (1 - alpha_next) / (1 - alpha)).sqrt()
+            c2 = ((1 - alpha_next) - torch.square(c1)).sqrt()
+
+            img = x_start * alpha_next.sqrt() + \
+                  c1 * torch.randn_like(img) + \
+                  c2 * pred_noise
+
+        img = unnormalize_to_zero_to_one(img)
         return img
 
     @torch.no_grad()
     def sample(self, batch_size = 16):
-        image_size = self.image_size
-        channels = self.channels
-        return self.p_sample_loop((batch_size, channels, image_size, image_size))
+        image_size, channels = self.image_size, self.channels
+        sample_fn = self.p_sample_loop if not self.is_ddim_sampling else self.ddim_sample
+        return sample_fn((batch_size, channels, image_size, image_size))
 
     @torch.no_grad()
     def interpolate(self, x1, x2, t = None, lam = 0.5):
@@ -420,43 +591,67 @@ 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_noisy = self.q_sample(x_start=x_start, t=t, noise=noise)
-        x_recon = self.denoise_fn(x_noisy, t)
+        x = self.q_sample(x_start = x_start, t = t, noise = noise)
+        model_out = self.model(x, t)
 
-        if self.loss_type == 'l1':
-            loss = (noise - x_recon).abs().mean()
-        elif self.loss_type == 'l2':
-            loss = F.mse_loss(noise, x_recon)
+        if self.objective == 'pred_noise':
+            target = noise
+        elif self.objective == 'pred_x0':
+            target = x_start
         else:
-            raise NotImplementedError()
+            raise ValueError(f'unknown objective {self.objective}')
 
-        return loss
+        loss = self.loss_fn(model_out, target, reduction = 'none')
+        loss = reduce(loss, 'b ... -> b (...)', 'mean')
 
-    def forward(self, x, *args, **kwargs):
-        b, c, h, w, device, img_size, = *x.shape, x.device, self.image_size
+        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
         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):
+class Dataset(Dataset):
+    def __init__(
+        self,
+        folder,
+        image_size,
+        exts = ['jpg', 'jpeg', 'png', 'tiff'],
+        augment_horizontal_flip = False,
+        convert_image_to = None
+    ):
         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),
-            transforms.RandomHorizontalFlip(),
-            transforms.CenterCrop(image_size),
-            transforms.ToTensor(),
-            transforms.Lambda(lambda t: (t * 2) - 1)
+        maybe_convert_fn = partial(convert_image_to, convert_image_to) if exists(convert_image_to) else nn.Identity()
+
+        self.transform = T.Compose([
+            T.Lambda(maybe_convert_fn),
+            T.Resize(image_size),
+            T.RandomHorizontalFlip() if augment_horizontal_flip else nn.Identity(),
+            T.CenterCrop(image_size),
+            T.ToTensor()
         ])
 
     def __len__(self):
@@ -475,89 +670,137 @@ class Trainer(object):
         diffusion_model,
         folder,
         *,
-        ema_decay = 0.995,
-        image_size = 128,
-        train_batch_size = 32,
-        train_lr = 2e-5,
+        train_batch_size = 16,
+        gradient_accumulate_every = 1,
+        augment_horizontal_flip = True,
+        train_lr = 1e-4,
         train_num_steps = 100000,
-        gradient_accumulate_every = 2,
+        ema_update_every = 10,
+        ema_decay = 0.995,
+        save_and_sample_every = 1000,
+        num_samples = 25,
+        results_folder = './results',
+        amp = False,
         fp16 = False,
-        step_start_ema = 2000
+        split_batches = True,
+        convert_image_to = None
     ):
         super().__init__()
+
+        self.accelerator = Accelerator(
+            split_batches = split_batches,
+            mixed_precision = 'fp16' if fp16 else 'no'
+        )
+
+        self.accelerator.native_amp = amp
+
         self.model = diffusion_model
-        self.ema = EMA(ema_decay)
-        self.ema_model = copy.deepcopy(self.model)
-        self.step_start_ema = step_start_ema
+
+        assert has_int_squareroot(num_samples), 'number of samples must have an integer square root'
+        self.num_samples = num_samples
+        self.save_and_sample_every = save_and_sample_every
 
         self.batch_size = train_batch_size
-        self.image_size = diffusion_model.image_size
         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.opt = Adam(diffusion_model.parameters(), lr=train_lr)
+        self.train_num_steps = train_num_steps
+        self.image_size = diffusion_model.image_size
+
+        # dataset and dataloader
+
+        self.ds = Dataset(folder, self.image_size, augment_horizontal_flip = augment_horizontal_flip, convert_image_to = convert_image_to)
+        dl = DataLoader(self.ds, batch_size = train_batch_size, shuffle = True, pin_memory = True, num_workers = cpu_count())
+
+        self.dl = cycle(dl)
+
+        # optimizer
+
+        self.opt = Adam(diffusion_model.parameters(), lr = train_lr)
+
+        # for logging results in a folder periodically
+
+        if self.accelerator.is_main_process:
+            self.ema = EMA(diffusion_model, beta = ema_decay, update_every = ema_update_every)
+
+            self.results_folder = Path(results_folder)
+            self.results_folder.mkdir(exist_ok = True)
+
+        # step counter state
 
         self.step = 0
 
-        assert not fp16 or fp16 and APEX_AVAILABLE, 'Apex must be installed in order for mixed precision training to be turned on'
+        # prepare model, dataloader, optimizer with accelerator
 
-        self.fp16 = fp16
-        if fp16:
-            (self.model, self.ema_model), self.opt = amp.initialize([self.model, self.ema_model], self.opt, opt_level='O1')
-
-        self.reset_parameters()
-
-    def reset_parameters(self):
-        self.ema_model.load_state_dict(self.model.state_dict())
-
-    def step_ema(self):
-        if self.step < self.step_start_ema:
-            self.reset_parameters()
-            return
-        self.ema.update_model_average(self.ema_model, self.model)
+        self.model, self.dl, self.opt = self.accelerator.prepare(self.model, self.dl, self.opt)
 
     def save(self, milestone):
+        if not self.accelerator.is_local_main_process:
+            return
+
         data = {
             'step': self.step,
-            'model': self.model.state_dict(),
-            'ema': self.ema_model.state_dict()
+            'model': self.accelerator.get_state_dict(self.model),
+            'opt': self.opt.state_dict(),
+            'ema': self.ema.state_dict(),
+            'scaler': self.accelerator.scaler.state_dict() if exists(self.accelerator.scaler) else None
         }
-        torch.save(data, str(RESULTS_FOLDER / f'model-{milestone}.pt'))
+
+        torch.save(data, str(self.results_folder / f'model-{milestone}.pt'))
 
     def load(self, milestone):
-        data = torch.load(str(RESULTS_FOLDER / f'model-{milestone}.pt'))
+        data = torch.load(str(self.results_folder / f'model-{milestone}.pt'))
+
+        model = self.accelerator.unwrap_model(self.model)
+        model.load_state_dict(data['model'])
 
         self.step = data['step']
-        self.model.load_state_dict(data['model'])
-        self.ema_model.load_state_dict(data['ema'])
+        self.opt.load_state_dict(data['opt'])
+        self.ema.load_state_dict(data['ema'])
+
+        if exists(self.accelerator.scaler) and exists(data['scaler']):
+            self.accelerator.scaler.load_state_dict(data['scaler'])
 
     def train(self):
-        backwards = partial(loss_backwards, self.fp16)
+        accelerator = self.accelerator
+        device = accelerator.device
 
-        while self.step < self.train_num_steps:
-            for i in range(self.gradient_accumulate_every):
-                data = next(self.dl).cuda()
-                loss = self.model(data)
-                print(f'{self.step}: {loss.item()}')
-                backwards(loss / self.gradient_accumulate_every, self.opt)
+        with tqdm(initial = self.step, total = self.train_num_steps, disable = not accelerator.is_main_process) as pbar:
 
-            self.opt.step()
-            self.opt.zero_grad()
+            while self.step < self.train_num_steps:
 
-            if self.step % UPDATE_EMA_EVERY == 0:
-                self.step_ema()
+                for _ in range(self.gradient_accumulate_every):
+                    data = next(self.dl).to(device)
 
-            if self.step != 0 and self.step % SAVE_AND_SAMPLE_EVERY == 0:
-                milestone = self.step // 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)
-                all_images = (all_images + 1) * 0.5
-                utils.save_image(all_images, str(RESULTS_FOLDER / f'sample-{milestone}.png'), nrow = 6)
-                self.save(milestone)
+                    with self.accelerator.autocast():
+                        loss = self.model(data)
+                        self.accelerator.backward(loss / self.gradient_accumulate_every)
 
-            self.step += 1
+                pbar.set_description(f'loss: {loss.item():.4f}')
 
-        print('training completed')
+                accelerator.wait_for_everyone()
+
+                self.opt.step()
+                self.opt.zero_grad()
+
+                accelerator.wait_for_everyone()
+
+                if accelerator.is_main_process:
+                    self.ema.to(device)
+                    self.ema.update()
+
+                    if self.step != 0 and self.step % self.save_and_sample_every == 0:
+                        self.ema.ema_model.eval()
+
+                        with torch.no_grad():
+                            milestone = self.step // self.save_and_sample_every
+                            batches = num_to_groups(self.num_samples, self.batch_size)
+                            all_images_list = list(map(lambda n: self.ema.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 = 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

@@ -0,0 +1,220 @@
+from math import sqrt
+import torch
+from torch import nn, einsum
+import torch.nn.functional as F
+
+from tqdm import tqdm
+from einops import rearrange, repeat, reduce
+
+# 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
+
+# tensor helpers
+
+def log(t, eps = 1e-20):
+    return torch.log(t.clamp(min = eps))
+
+# 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
+
+# main class
+
+class ElucidatedDiffusion(nn.Module):
+    def __init__(
+        self,
+        net,
+        *,
+        image_size,
+        channels = 3,
+        num_sample_steps = 32, # number of sampling steps
+        sigma_min = 0.002,     # min noise level
+        sigma_max = 80,        # max noise level
+        sigma_data = 0.5,      # standard deviation of data distribution
+        rho = 7,               # controls the sampling schedule
+        P_mean = -1.2,         # mean of log-normal distribution from which noise is drawn for training
+        P_std = 1.2,           # standard deviation of log-normal distribution from which noise is drawn for training
+        S_churn = 80,          # parameters for stochastic sampling - depends on dataset, Table 5 in apper
+        S_tmin = 0.05,
+        S_tmax = 50,
+        S_noise = 1.003,
+    ):
+        super().__init__()
+        assert net.learned_sinusoidal_cond
+
+        self.net = net
+
+        # image dimensions
+
+        self.channels = channels
+        self.image_size = image_size
+
+        # parameters
+
+        self.sigma_min = sigma_min
+        self.sigma_max = sigma_max
+        self.sigma_data = sigma_data
+
+        self.rho = rho
+
+        self.P_mean = P_mean
+        self.P_std = P_std
+
+        self.num_sample_steps = num_sample_steps  # otherwise known as N in the paper
+
+        self.S_churn = S_churn
+        self.S_tmin = S_tmin
+        self.S_tmax = S_tmax
+        self.S_noise = S_noise
+
+    @property
+    def device(self):
+        return next(self.net.parameters()).device
+
+    # derived preconditioning params - Table 1
+
+    def c_skip(self, sigma):
+        return (self.sigma_data ** 2) / (sigma ** 2 + self.sigma_data ** 2)
+
+    def c_out(self, sigma):
+        return sigma * self.sigma_data * (self.sigma_data ** 2 + sigma ** 2) ** -0.5
+
+    def c_in(self, sigma):
+        return 1 * (sigma ** 2 + self.sigma_data ** 2) ** -0.5
+
+    def c_noise(self, sigma):
+        return log(sigma) * 0.25
+
+    # preconditioned network output
+    # equation (7) in the paper
+
+    def preconditioned_network_forward(self, noised_images, sigma, clamp = False):
+        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)
+        )
+
+        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
+
+    def sample_schedule(self, num_sample_steps = None):
+        num_sample_steps = default(num_sample_steps, self.num_sample_steps)
+
+        N = num_sample_steps
+        inv_rho = 1 / self.rho
+
+        steps = torch.arange(num_sample_steps, device = self.device, dtype = torch.float32)
+        sigmas = (self.sigma_max ** inv_rho + steps / (N - 1) * (self.sigma_min ** inv_rho - self.sigma_max ** inv_rho)) ** self.rho
+
+        sigmas = F.pad(sigmas, (0, 1), value = 0.) # last step is sigma value of 0.
+        return sigmas
+
+    @torch.no_grad()
+    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)
+
+        # get the schedule, which is returned as (sigma, gamma) tuple, and pair up with the next sigma and gamma
+
+        sigmas = self.sample_schedule(num_sample_steps)
+
+        gammas = torch.where(
+            (sigmas >= self.S_tmin) & (sigmas <= self.S_tmax),
+            min(self.S_churn / num_sample_steps, sqrt(2) - 1),
+            0.
+        )
+
+        sigmas_and_gammas = list(zip(sigmas[:-1], sigmas[1:], gammas[:-1]))
+
+        # images is noise at the beginning
+
+        init_sigma = sigmas[0]
+
+        images = init_sigma * torch.randn(shape, device = self.device)
+
+        # gradually denoise
+
+        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 = 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, clamp = clamp)
+            denoised_over_sigma = (images_hat - model_output) / sigma_hat
+
+            images_next = images_hat + (sigma_next - sigma_hat) * denoised_over_sigma
+
+            # 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)
+                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
+
+        images = images.clamp(-1., 1.)
+        return unnormalize_to_zero_to_one(images)
+
+    # 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
+
+        assert h == image_size and w == image_size, f'height and width of image must be {image_size}'
+        assert c == channels, 'mismatch of image channels'
+
+        images = normalize_to_neg_one_to_one(images)
+
+        sigmas = self.noise_distribution(batch_size)
+        padded_sigmas = rearrange(sigmas, 'b -> b 1 1 1')
+
+        noise = torch.randn_like(images)
+
+        noised_images = images + padded_sigmas * noise  # alphas are 1. in the paper
+
+        denoised = self.preconditioned_network_forward(noised_images, sigmas)
+
+        losses = F.mse_loss(denoised, images, reduction = 'none')
+        losses = reduce(losses, 'b ... -> b', 'mean')
+
+        losses = losses * self.loss_weight(sigmas)
+
+        return losses.mean()

denoising_diffusion_pytorch/learned_gaussian_diffusion.py

@@ -0,0 +1,149 @@
+import torch
+from collections import namedtuple
+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)
+
+ModelPrediction = namedtuple('ModelPrediction', ['pred_noise', 'pred_x_start', 'pred_variance'])
+
+# 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-15):
+    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,
+        model,
+        vb_loss_weight = 0.001,  # lambda was 0.001 in the paper
+        *args,
+        **kwargs
+    ):
+        super().__init__(model, *args, **kwargs)
+        assert model.out_dim == (model.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 model_predictions(self, x, t):
+        model_output = self.model(x, t)
+        model_output, pred_variance = model_output.chunk(2, dim = 1)
+
+        if self.objective == 'pred_noise':
+            pred_noise = model_output
+            x_start = self.predict_start_from_noise(x, t, model_output)
+
+        elif self.objective == 'pred_x0':
+            pred_noise = self.predict_noise_from_start(x, t, model_output)
+            x_start = model_output
+
+        return ModelPrediction(pred_noise, x_start, pred_variance)
+
+    def p_mean_variance(self, *, x, t, clip_denoised, model_output = None):
+        model_output = default(model_output, lambda: self.model(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.model(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

@@ -0,0 +1,81 @@
+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,
+        model,
+        *args,
+        pred_noise_loss_weight = 0.1,
+        pred_x_start_loss_weight = 0.1,
+        **kwargs
+    ):
+        super().__init__(model, *args, **kwargs)
+        channels = model.channels
+        assert model.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'
+        assert not self.is_ddim_sampling, 'ddim sampling cannot be used'
+
+        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.model(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.model(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,21 @@ from setuptools import setup, find_packages
 setup(
   name = 'denoising-diffusion-pytorch',
   packages = find_packages(),
-  version = '0.6.5',
+  version = '0.25.0',
   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=[
+    'accelerate',
     'einops',
-    'numpy',
+    'ema-pytorch',
     'pillow',
     'torch',
     'torchvision',