0.27.6

Update step before training checkpoint

.gitignore

@@ -1,3 +1,6 @@
+# Generation results
+results/
+
 # Byte-compiled / optimized / DLL files
 __pycache__/
 *.py[cod]

README.md

@@ -1,6 +1,20 @@
-## Denoising Diffusion Probabilistic Model, in Pytorch (wip)
+<img src="./images/denoising-diffusion.png" width="500px"></img>
 
-Implementation of <a href="https://arxiv.org/abs/2006.11239">Denoising Diffusion Probabilistic Model</a> in Pytorch
+## 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>
+
+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>
+
+Update: Turns out none of the technicalities really matters at all | <a href="https://arxiv.org/abs/2208.09392">"Cold Diffusion" paper</a>
+
+<img src="./images/sample.png" width="500px"><img>
+
+[![PyPI version](https://badge.fury.io/py/denoising-diffusion-pytorch.svg)](https://badge.fury.io/py/denoising-diffusion-pytorch)
 
 ## Install
 
@@ -21,30 +35,161 @@ model = Unet(
 
 diffusion = GaussianDiffusion(
     model,
-    beta_start = 0.0001,
-    beta_end = 0.02,
-    num_diffusion_timesteps = 1000,   # number of steps
-    loss_type = 'l1'                  # L1 or L2
+    image_size = 128,
+    timesteps = 1000,   # number of steps
+    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
 
-sampled_images = diffusion.p_sample_loop((1, 3, 128, 128))
-sampled_images.shape # (1, 3, 128, 128)
+sampled_images = diffusion.sample(batch_size = 4)
+sampled_images.shape # (4, 3, 128, 128)
+```
+
+Or, if you simply want to pass in a folder name and the desired image dimensions, you can use the `Trainer` class to easily train a model.
+
+```python
+from denoising_diffusion_pytorch import Unet, GaussianDiffusion, Trainer
+
+model = Unet(
+    dim = 64,
+    dim_mults = (1, 2, 4, 8)
+).cuda()
+
+diffusion = GaussianDiffusion(
+    model,
+    image_size = 128,
+    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 = 8e-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
+)
+
+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{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},
+    volume  = {abs/2010.02502}
+}
+```
+
+```bibtex
+@misc{chen2022analog,
+    title   = {Analog Bits: Generating Discrete Data using Diffusion Models with Self-Conditioning},
+    author  = {Ting Chen and Ruixiang Zhang and Geoffrey Hinton},
+    year    = {2022},
+    eprint  = {2208.04202},
+    archivePrefix = {arXiv},
+    primaryClass = {cs.CV}
+}
+```
+
+```bibtex
+@article{Qiao2019WeightS,
+    title   = {Weight Standardization},
+    author  = {Siyuan Qiao and Huiyu Wang and Chenxi Liu and Wei Shen and Alan Loddon Yuille},
+    journal = {ArXiv},
+    year    = {2019},
+    volume  = {abs/1903.10520}
 }
 ```

denoising_diffusion_pytorch/__init__.py

@@ -1 +1,6 @@
-from denoising_diffusion_pytorch.denoising_diffusion_pytorch import GaussianDiffusion, Unet
+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,288 @@
+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
+        assert not model.self_condition, 'not supported yet'
+
+        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

@@ -1,13 +1,31 @@
 import math
-import torch
-from inspect import isfunction
+import copy
+from pathlib import Path
+from random import random
 from functools import partial
+from collections import namedtuple
+from multiprocessing import cpu_count
+
+import torch
 from torch import nn, einsum
 import torch.nn.functional as F
+from torch.utils.data import Dataset, DataLoader
 
-import numpy as np
-from tqdm import tqdm
-from einops import rearrange
+from torch.optim import Adam
+from torchvision import transforms as T, utils
+
+from einops import rearrange, reduce
+from einops.layers.torch import Rearrange
+
+from PIL import Image
+from tqdm.auto import tqdm
+from ema_pytorch import EMA
+
+from accelerate import Accelerator
+
+# constants
+
+ModelPrediction =  namedtuple('ModelPrediction', ['pred_noise', 'pred_x_start'])
 
 # helpers functions
 
@@ -17,10 +35,39 @@ def exists(x):
 def default(val, d):
     if exists(val):
         return val
-    return d() if isfunction(d) else d
+    return d() if callable(d) else d
 
-def normal_kl(mean1, logvar1, mean2, logvar2):
-    return 0.5 * (-1. + logvar2 - logvar1 + torch.exp(logvar1 - logvar2) + torch.exp(-logvar2) * (mean1 - mean2) ** 2)
+def cycle(dl):
+    while True:
+        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
+    arr = [divisor] * groups
+    if remainder > 0:
+        arr.append(remainder)
+    return arr
+
+def convert_image_to(img_type, image):
+    if image.mode != img_type:
+        return image.convert(img_type)
+    return image
+
+def l2norm(t):
+    return F.normalize(t, dim = -1)
+
+# 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
 
@@ -32,6 +79,53 @@ 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 WeightStandardizedConv2d(nn.Conv2d):
+    """
+    https://arxiv.org/abs/1903.10520
+    weight standardization purportedly works synergistically with group normalization
+    """
+    def forward(self, x):
+        eps = 1e-5 if x.dtype == torch.float32 else 1e-3
+
+        weight = self.weight
+        mean = reduce(weight, 'o ... -> o 1 1 1', 'mean')
+        var = reduce(weight, 'o ... -> o 1 1 1', partial(torch.var, unbiased = False))
+        normalized_weight = (weight - mean) * (var + eps).rsqrt()
+
+        return F.conv2d(x, normalized_weight, self.bias, self.stride, self.padding, self.dilation, self.groups)
+
+class LayerNorm(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        self.g = nn.Parameter(torch.ones(1, dim, 1, 1))
+
+    def forward(self, x):
+        eps = 1e-5 if x.dtype == torch.float32 else 1e-3
+        var = torch.var(x, dim = 1, unbiased = False, keepdim = True)
+        mean = torch.mean(x, dim = 1, keepdim = True)
+        return (x - mean) * (var + eps).rsqrt() * self.g
+
+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__()
@@ -46,99 +140,175 @@ 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, dim):
-        super().__init__()
-        self.g = nn.Parameter(torch.zeros(1))
-
-    def forward(self, x):
-        return 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 = 32):
+    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 = WeightStandardizedConv2d(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 = 32):
+    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 = 8, dim_head = 32):
+    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, 1, bias = False)
+        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
+        v = v / (h * w)
+
+        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, scale = 10):
+        super().__init__()
+        self.scale = scale
+        self.heads = heads
+        hidden_dim = dim_head * heads
+        self.to_qkv = nn.Conv2d(dim, hidden_dim * 3, 1, bias = False)
         self.to_out = nn.Conv2d(hidden_dim, dim, 1)
 
     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)
-        q = q.softmax(dim=-2)
-        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, k = map(l2norm, (q, k))
+
+        sim = einsum('b h d i, b h d j -> b h i j', q, k) * self.scale
+        attn = sim.softmax(dim = -1)
+        out = einsum('b h i j, b h d j -> b h i d', attn, v)
+        out = rearrange(out, 'b h (x y) d -> b (h d) x y', x = h, y = w)
         return self.to_out(out)
 
 # model
 
 class Unet(nn.Module):
-    def __init__(self, dim, out_dim = None, dim_mults=(1, 2, 4, 8), groups = 32):
+    def __init__(
+        self,
+        dim,
+        init_dim = None,
+        out_dim = None,
+        dim_mults=(1, 2, 4, 8),
+        channels = 3,
+        self_condition = False,
+        resnet_block_groups = 8,
+        learned_variance = False,
+        learned_sinusoidal_cond = False,
+        learned_sinusoidal_dim = 16
+    ):
         super().__init__()
-        dims = [3, *map(lambda m: dim * m, dim_mults)]
+
+        # determine dimensions
+
+        self.channels = channels
+        self.self_condition = self_condition
+        input_channels = channels * (2 if self_condition else 1)
+
+        init_dim = default(init_dim, dim)
+        self.init_conv = nn.Conv2d(input_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)
@@ -147,53 +317,72 @@ class Unet(nn.Module):
             is_last = ind >= (num_resolutions - 1)
 
             self.downs.append(nn.ModuleList([
-                ResnetBlock(dim_in, 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),
-                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, 3)
-        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)
 
-    def forward(self, x, time):
-        t = self.time_pos_emb(time)
-        t = self.mlp(t)
+        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, x_self_cond = None):
+        if self.self_condition:
+            x_self_cond = default(x_self_cond, lambda: torch.zeros_like(x))
+            x = torch.cat((x_self_cond, x), dim = 1)
+
+        x = self.init_conv(x)
+        r = x.clone()
+
+        t = self.time_mlp(time)
 
         h = []
 
-        for resnet, attn, downsample in self.downs:
-            x = resnet(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, attn, upsample in self.ups:
-            x = torch.cat((x, h.pop()), dim=1)
-            x = resnet(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
@@ -203,54 +392,109 @@ 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):
+    """
+    cosine schedule
+    as proposed in https://openreview.net/forum?id=-NEXDKk8gZ
+    """
+    steps = timesteps + 1
+    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 torch.clip(betas, 0, 0.999)
 
 class GaussianDiffusion(nn.Module):
-    def __init__(self, denoise_fn, beta_start=0.0001, beta_end=0.02, num_diffusion_timesteps=1000, loss_type='l1'):
+    def __init__(
+        self,
+        model,
+        *,
+        image_size,
+        timesteps = 1000,
+        sampling_timesteps = None,
+        loss_type = 'l1',
+        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__()
-        self.denoise_fn = denoise_fn
+        assert not (type(self) == GaussianDiffusion and model.channels != model.out_dim)
+        assert not model.learned_sinusoidal_cond
+
+        self.model = model
+        self.channels = self.model.channels
+        self.self_condition = self.model.self_condition
+
+        self.image_size = image_size
+
+        self.objective = objective
+
+        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 = torch.cumprod(alphas, axis=0)
+        alphas_cumprod_prev = F.pad(alphas_cumprod[:-1], (1, 0), value = 1.)
 
-        self.np_betas = betas = np.linspace(beta_start, beta_end, num_diffusion_timesteps).astype(np.float64)
         timesteps, = betas.shape
         self.num_timesteps = int(timesteps)
         self.loss_type = loss_type
 
-        alphas = 1. - betas
-        alphas_cumprod = np.cumprod(alphas, axis=0)
-        alphas_cumprod_prev = np.append(1., alphas_cumprod[:-1])
+        # sampling related parameters
 
-        to_torch = partial(torch.tensor, dtype=torch.float32)
+        self.sampling_timesteps = default(sampling_timesteps, timesteps) # default num sampling timesteps to number of timesteps at training
 
-        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))
+        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 (
@@ -258,6 +502,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 (
+            (extract(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t - x0) / \
+            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 +
@@ -267,33 +517,110 @@ 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 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))
+    def model_predictions(self, x, t, x_self_cond = None):
+        model_output = self.model(x, t, x_self_cond)
+
+        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, x_self_cond = None, clip_denoised = True):
+        preds = self.model_predictions(x, t, x_self_cond)
+        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)
-        return model_mean, posterior_variance, posterior_log_variance
+        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, x_start
 
     @torch.no_grad()
-    def p_sample(self, x, t, clip_denoised=True, repeat_noise=False):
+    def p_sample(self, x, t: int, x_self_cond = None, 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, x_start = self.p_mean_variance(x = x, t = batched_times, x_self_cond = x_self_cond, clip_denoised = clip_denoised)
+        noise = torch.randn_like(x) if t > 0 else 0. # no noise if t == 0
+        pred_img = model_mean + (0.5 * model_log_variance).exp() * noise
+        return pred_img, x_start
 
     @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):
+        x_start = None
+
+        for t in tqdm(reversed(range(0, self.num_timesteps)), desc = 'sampling loop time step', total = self.num_timesteps):
+            self_cond = x_start if self.self_condition else None
+            img, x_start = self.p_sample(img, t, self_cond)
+
+        img = unnormalize_to_zero_to_one(img)
+        return img
+
+    @torch.no_grad()
+    def ddim_sample(self, shape, clip_denoised = True):
+        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(filter(lambda a: a[0] > a[1], zip(times[:-1], times[1:])))
+
+        img = torch.randn(shape, device = device)
+
+        x_start = None
+
+        for time, time_next in tqdm(time_pairs, desc = 'sampling loop time step'):
+            alpha = self.alphas_cumprod[time]
+            alpha_next = self.alphas_cumprod[time_next]
+
+            time_cond = torch.full((batch,), time, device = device, dtype = torch.long)
+
+            self_cond = x_start if self.self_condition else None
+
+            pred_noise, x_start, *_ = self.model_predictions(img, time_cond, self_cond)
+
+            if clip_denoised:
+                x_start.clamp_(-1., 1.)
+
+            sigma = eta * ((1 - alpha / alpha_next) * (1 - alpha_next) / (1 - alpha)).sqrt()
+            c = ((1 - alpha_next) - sigma ** 2).sqrt()
+
+            noise = torch.randn_like(img) if time_next > 0 else 0.
+
+            img = x_start * alpha_next.sqrt() + \
+                  c * pred_noise + \
+                  sigma * noise
+
+        img = unnormalize_to_zero_to_one(img)
+        return img
+
+    @torch.no_grad()
+    def sample(self, batch_size = 16):
+        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):
+        b, *_, device = *x1.shape, x1.device
+        t = default(t, self.num_timesteps - 1)
+
+        assert x1.shape == x2.shape
+
+        t_batched = torch.stack([torch.tensor(t, device = device)] * b)
+        xt1, xt2 = map(lambda x: self.q_sample(x, t = t_batched), (x1, x2))
+
+        img = (1 - lam) * xt1 + lam * xt2
+        for i in tqdm(reversed(range(0, t)), desc = 'interpolation sample time step', total = t):
             img = self.p_sample(img, torch.full((b,), i, device=device, dtype=torch.long))
+
         return img
 
     def q_sample(self, x_start, t, noise=None):
@@ -304,23 +631,238 @@ 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)
+        # noise sample
 
-        if self.loss_type == 'l1':
-            loss = (noise - x_recon).abs().mean()
-        elif self.loss_type == 'l2':
-            loss = F.mse_loss(noise, x_recon)
+        x = self.q_sample(x_start = x_start, t = t, noise = noise)
+
+        # if doing self-conditioning, 50% of the time, predict x_start from current set of times
+        # and condition with unet with that
+        # this technique will slow down training by 25%, but seems to lower FID significantly
+
+        x_self_cond = None
+        if self.self_condition and random() < 0.5:
+            with torch.no_grad():
+                x_self_cond = self.model_predictions(x, t).pred_x_start
+                x_self_cond.detach_()
+
+        # predict and take gradient step
+
+        model_out = self.model(x, t, x_self_cond)
+
+        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, *_, device = *x.shape, x.device
-        t = torch.randint(0, 1000, (b,), device=device).long()
-        return self.p_losses(x, t, *args, **kwargs)
+        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()
+
+        img = normalize_to_neg_one_to_one(img)
+        return self.p_losses(img, t, *args, **kwargs)
+
+# dataset classes
+
+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}')]
+
+        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):
+        return len(self.paths)
+
+    def __getitem__(self, index):
+        path = self.paths[index]
+        img = Image.open(path)
+        return self.transform(img)
+
+# trainer class
+
+class Trainer(object):
+    def __init__(
+        self,
+        diffusion_model,
+        folder,
+        *,
+        train_batch_size = 16,
+        gradient_accumulate_every = 1,
+        augment_horizontal_flip = True,
+        train_lr = 1e-4,
+        train_num_steps = 100000,
+        ema_update_every = 10,
+        ema_decay = 0.995,
+        adam_betas = (0.9, 0.99),
+        save_and_sample_every = 1000,
+        num_samples = 25,
+        results_folder = './results',
+        amp = False,
+        fp16 = False,
+        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
+
+        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.gradient_accumulate_every = gradient_accumulate_every
+
+        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())
+
+        dl = self.accelerator.prepare(dl)
+        self.dl = cycle(dl)
+
+        # optimizer
+
+        self.opt = Adam(diffusion_model.parameters(), lr = train_lr, betas = adam_betas)
+
+        # 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
+
+        # prepare model, dataloader, optimizer with accelerator
+
+        self.model, self.opt = self.accelerator.prepare(self.model, self.opt)
+
+    def save(self, milestone):
+        if not self.accelerator.is_local_main_process:
+            return
+
+        data = {
+            'step': self.step,
+            '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(self.results_folder / f'model-{milestone}.pt'))
+
+    def load(self, milestone):
+        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.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):
+        accelerator = self.accelerator
+        device = accelerator.device
+
+        with tqdm(initial = self.step, total = self.train_num_steps, disable = not accelerator.is_main_process) as pbar:
+
+            while self.step < self.train_num_steps:
+
+                total_loss = 0.
+
+                for _ in range(self.gradient_accumulate_every):
+                    data = next(self.dl).to(device)
+
+                    with self.accelerator.autocast():
+                        loss = self.model(data)
+                        loss = loss / self.gradient_accumulate_every
+                        total_loss += loss.item()
+
+                    self.accelerator.backward(loss)
+
+                pbar.set_description(f'loss: {total_loss:.4f}')
+
+                accelerator.wait_for_everyone()
+
+                self.opt.step()
+                self.opt.zero_grad()
+
+                accelerator.wait_for_everyone()
+
+                self.step += 1
+                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)
+
+                pbar.update(1)
+
+        accelerator.print('training complete')

denoising_diffusion_pytorch/elucidated_diffusion.py

@@ -0,0 +1,240 @@
+from math import sqrt
+from random import random
+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.self_condition = net.self_condition
+
+        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, self_cond = None, 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),
+            self_cond
+        )
+
+        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)
+
+        # for self conditioning
+
+        x_start = None
+
+        # 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
+
+            self_cond = x_start if self.self_condition else None
+
+            model_output = self.preconditioned_network_forward(images_hat, sigma_hat, self_cond, clamp = clamp)
+            denoised_over_sigma = (images_hat - model_output) / sigma_hat
+
+            images_next = images_hat + (sigma_next - sigma_hat) * denoised_over_sigma
+
+            # second order correction, if not the last timestep
+
+            if sigma_next != 0:
+                self_cond = model_output if self.self_condition else None
+
+                model_output_next = self.preconditioned_network_forward(images_next, sigma_next, self_cond, clamp = clamp)
+                denoised_prime_over_sigma = (images_next - model_output_next) / sigma_next
+                images_next = images_hat + 0.5 * (sigma_next - sigma_hat) * (denoised_over_sigma + denoised_prime_over_sigma)
+
+            images = images_next
+            x_start = model_output
+
+        images = images.clamp(-1., 1.)
+        return unnormalize_to_zero_to_one(images)
+
+    # 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
+
+        self_cond = None
+
+        if self.self_condition and random() < 0.5:
+            # from hinton's group's bit diffusion paper
+            with torch.no_grad():
+                self_cond = self.preconditioned_network_forward(noised_images, sigmas)
+                self_cond.detach_()
+
+        denoised = self.preconditioned_network_forward(noised_images, sigmas, self_cond)
+
+        losses = F.mse_loss(denoised, images, reduction = 'none')
+        losses = reduce(losses, 'b ... -> b', 'mean')
+
+        losses = losses * self.loss_weight(sigmas)
+
+        return losses.mean()

denoising_diffusion_pytorch/learned_gaussian_diffusion.py

@@ -0,0 +1,151 @@
+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`'
+        assert not model.self_condition, 'not supported yet'
+
+        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,82 @@
+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 model.self_condition, 'not supported yet'
+        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,20 +3,24 @@ from setuptools import setup, find_packages
 setup(
   name = 'denoising-diffusion-pytorch',
   packages = find_packages(),
-  version = '0.0.1',
+  version = '0.27.6',
   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',
     'tqdm'
   ],
   classifiers=[
@@ -26,4 +30,4 @@ setup(
     'License :: OSI Approved :: MIT License',
     'Programming Language :: Python :: 3.6',
   ],
-)
+)