From 09b8a1c8050928eb7b8059eb224ac4d9d3b1a94a Mon Sep 17 00:00:00 2001
From: Phil Wang <lucidrains@gmail.com>
Date: Tue, 28 Jun 2022 13:11:56 -0700
Subject: [PATCH 01/10] some basic scaffold for elucidating diffusion and
 derived values

---
 README.md                                     |  10 ++
 denoising_diffusion_pytorch/__init__.py       |   1 +
 .../elucidated_diffusion.py                   | 156 ++++++++++++++++++
 3 files changed, 167 insertions(+)
 create mode 100644 denoising_diffusion_pytorch/elucidated_diffusion.py

diff --git a/README.md b/README.md
index bc2a3c5..906e46e 100644
--- a/README.md
+++ b/README.md
@@ -133,3 +133,13 @@ Samples and model checkpoints will be logged to `./results` periodically
     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}
+}
+```
diff --git a/denoising_diffusion_pytorch/__init__.py b/denoising_diffusion_pytorch/__init__.py
index 75ab437..8108a41 100644
--- a/denoising_diffusion_pytorch/__init__.py
+++ b/denoising_diffusion_pytorch/__init__.py
@@ -3,3 +3,4 @@ from denoising_diffusion_pytorch.denoising_diffusion_pytorch import GaussianDiff
 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
diff --git a/denoising_diffusion_pytorch/elucidated_diffusion.py b/denoising_diffusion_pytorch/elucidated_diffusion.py
new file mode 100644
index 0000000..a7d906b
--- /dev/null
+++ b/denoising_diffusion_pytorch/elucidated_diffusion.py
@@ -0,0 +1,156 @@
+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,
+        denoise_fn,
+        *,
+        image_size,
+        channels = 3,
+        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 denoise_fn.learned_sinusoidal_cond
+
+        self.denoise_fn = denoise_fn
+
+        # 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.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.denoise_fn.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 (sigma ** 2 + self.sigma_data ** 2) * -0.5
+
+    def c_noise(self, sigma):
+        """ apparently empirically derived """
+        return log(sigma) ** 0.25
+
+    # noise distribution
+
+    def noise_distribution(self, batch_size):
+        return (self.P_mean + self.P_std * torch.randn((batch_size,), device = self.device)).exp()
+
+    def loss_weight(self, sigma):
+        return (sigma ** 2 + self.sigma_data ** 2) * (sigma * self.sigma_data) ** -2
+
+    # sampling related functions
+
+    @torch.no_grad()
+    def sample_one_timestep(self, x, time, time_next):
+        batch, *_, device = *x.shape, x.device
+        return x
+
+    @torch.no_grad()
+    def sample_all_timesteps(self, shape):
+        batch = shape[0]
+
+        img = torch.randn(shape, device = self.device)
+        steps = torch.linspace(1., 0., 100 + 1, device = self.device)
+
+        for i in tqdm(range(100), desc = 'sampling loop time step', total = 100):
+            times = steps[i]
+            times_next = steps[i + 1]
+            img = self.sample_one_timestep(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.sample_all_timesteps((batch_size, self.channels, self.image_size, self.image_size))
+
+    # training related functions - noise prediction
+
+    def add_noise(self, x_start, times, noise = None):
+        noise = default(noise, lambda: torch.randn_like(x_start))
+        x_noised =  x_start + noise
+        return x_noised, noise.mean(dim = (1, 2, 3))
+
+    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 forward(self, images):
+        b, c, h, w, device, image_size, = *images.shape, images.device, self.image_size
+        assert h == image_size and w == image_size, f'height and width of image must be {image_size}'
+
+        times = self.random_times(b)
+        images = normalize_to_neg_one_to_one(images)
+
+        noise = torch.randn_like(images)
+
+        noise_images, log_snr = self.add_noise(x_start = images, times = times, noise = noise)
+        model_out = self.denoise_fn(noise_images, log_snr)
+
+        losses = F.mse_loss(model_out, noise, reduction = 'none')
+        losses = reduce(losses, 'b ... -> b', 'mean')
+        return losses.mean()

From 06b2e52645bc7fb458f2b7efe954bf8416467a4e Mon Sep 17 00:00:00 2001
From: Phil Wang <lucidrains@gmail.com>
Date: Tue, 28 Jun 2022 13:35:06 -0700
Subject: [PATCH 02/10] get training working

---
 .../elucidated_diffusion.py                   | 55 ++++++++++++-------
 1 file changed, 36 insertions(+), 19 deletions(-)

diff --git a/denoising_diffusion_pytorch/elucidated_diffusion.py b/denoising_diffusion_pytorch/elucidated_diffusion.py
index a7d906b..6e938ec 100644
--- a/denoising_diffusion_pytorch/elucidated_diffusion.py
+++ b/denoising_diffusion_pytorch/elucidated_diffusion.py
@@ -33,7 +33,7 @@ def unnormalize_to_zero_to_one(t):
 class ElucidatedDiffusion(nn.Module):
     def __init__(
         self,
-        denoise_fn,
+        net,
         *,
         image_size,
         channels = 3,
@@ -49,9 +49,9 @@ class ElucidatedDiffusion(nn.Module):
         S_noise = 1.003
     ):
         super().__init__()
-        assert denoise_fn.learned_sinusoidal_cond
+        assert net.learned_sinusoidal_cond
 
-        self.denoise_fn = denoise_fn
+        self.net = net
 
         # image dimensions
 
@@ -76,7 +76,7 @@ class ElucidatedDiffusion(nn.Module):
 
     @property
     def device(self):
-        return next(self.denoise_fn.parameters()).device
+        return next(self.net.parameters()).device
 
     # derived preconditioning params - Table 1
 
@@ -91,7 +91,7 @@ class ElucidatedDiffusion(nn.Module):
 
     def c_noise(self, sigma):
         """ apparently empirically derived """
-        return log(sigma) ** 0.25
+        return log(sigma) * 0.25
 
     # noise distribution
 
@@ -101,7 +101,20 @@ class ElucidatedDiffusion(nn.Module):
     def loss_weight(self, sigma):
         return (sigma ** 2 + self.sigma_data ** 2) * (sigma * self.sigma_data) ** -2
 
-    # sampling related functions
+    # preconditioned network output
+    # equation (7) in the paper
+
+    def preconditioned_network_forward(self, noised_images, sigma):
+        padded_sigma = rearrange(sigma, 'b -> b 1 1 1')
+
+        net_out = self.net(
+            self.c_in(padded_sigma) * noised_images,
+            self.c_noise(sigma)
+        )
+
+        return self.c_skip(padded_sigma) * noised_images +  self.c_out(padded_sigma) * net_out
+
+    # sampling
 
     @torch.no_grad()
     def sample_one_timestep(self, x, time, time_next):
@@ -110,25 +123,21 @@ class ElucidatedDiffusion(nn.Module):
 
     @torch.no_grad()
     def sample_all_timesteps(self, shape):
-        batch = shape[0]
-
-        img = torch.randn(shape, device = self.device)
+        images = torch.randn(shape, device = self.device)
         steps = torch.linspace(1., 0., 100 + 1, device = self.device)
 
         for i in tqdm(range(100), desc = 'sampling loop time step', total = 100):
             times = steps[i]
             times_next = steps[i + 1]
-            img = self.sample_one_timestep(img, times, times_next)
+            images = self.sample_one_timestep(images, times, times_next)
 
-        img.clamp_(-1., 1.)
-        img = unnormalize_to_zero_to_one(img)
-        return img
+        return unnormalize_to_zero_to_one(images)
 
     @torch.no_grad()
     def sample(self, batch_size = 16):
         return self.sample_all_timesteps((batch_size, self.channels, self.image_size, self.image_size))
 
-    # training related functions - noise prediction
+    # training
 
     def add_noise(self, x_start, times, noise = None):
         noise = default(noise, lambda: torch.randn_like(x_start))
@@ -140,17 +149,25 @@ class ElucidatedDiffusion(nn.Module):
         return torch.zeros((batch_size,), device = self.device).float().uniform_(0, 1)
 
     def forward(self, images):
-        b, c, h, w, device, image_size, = *images.shape, images.device, self.image_size
+        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'
 
-        times = self.random_times(b)
         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)
 
-        noise_images, log_snr = self.add_noise(x_start = images, times = times, noise = noise)
-        model_out = self.denoise_fn(noise_images, log_snr)
+        noised_images = images + padded_sigmas * noise  # alphas are 1. in the paper
 
-        losses = F.mse_loss(model_out, noise, reduction = 'none')
+        model_out = self.preconditioned_network_forward(noised_images, sigmas)
+
+        losses = F.mse_loss(model_out, images, reduction = 'none')
         losses = reduce(losses, 'b ... -> b', 'mean')
+
+        losses = losses * self.loss_weight(sigmas)
+
         return losses.mean()

From c3d160701958391d64eec55d3474e2ac61378845 Mon Sep 17 00:00:00 2001
From: Phil Wang <lucidrains@gmail.com>
Date: Tue, 28 Jun 2022 13:38:37 -0700
Subject: [PATCH 03/10] cleanup

---
 denoising_diffusion_pytorch/elucidated_diffusion.py | 9 ---------
 1 file changed, 9 deletions(-)

diff --git a/denoising_diffusion_pytorch/elucidated_diffusion.py b/denoising_diffusion_pytorch/elucidated_diffusion.py
index 6e938ec..002e87d 100644
--- a/denoising_diffusion_pytorch/elucidated_diffusion.py
+++ b/denoising_diffusion_pytorch/elucidated_diffusion.py
@@ -139,15 +139,6 @@ class ElucidatedDiffusion(nn.Module):
 
     # training
 
-    def add_noise(self, x_start, times, noise = None):
-        noise = default(noise, lambda: torch.randn_like(x_start))
-        x_noised =  x_start + noise
-        return x_noised, noise.mean(dim = (1, 2, 3))
-
-    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 forward(self, images):
         batch_size, c, h, w, device, image_size, channels = *images.shape, images.device, self.image_size, self.channels
 

From be2bd8d32031800831a3e6bf15cb7f8cb2a65a94 Mon Sep 17 00:00:00 2001
From: Phil Wang <lucidrains@gmail.com>
Date: Tue, 28 Jun 2022 13:42:41 -0700
Subject: [PATCH 04/10] cleanup again

---
 denoising_diffusion_pytorch/elucidated_diffusion.py | 13 +++----------
 1 file changed, 3 insertions(+), 10 deletions(-)

diff --git a/denoising_diffusion_pytorch/elucidated_diffusion.py b/denoising_diffusion_pytorch/elucidated_diffusion.py
index 002e87d..1dfff27 100644
--- a/denoising_diffusion_pytorch/elucidated_diffusion.py
+++ b/denoising_diffusion_pytorch/elucidated_diffusion.py
@@ -117,26 +117,19 @@ class ElucidatedDiffusion(nn.Module):
     # sampling
 
     @torch.no_grad()
-    def sample_one_timestep(self, x, time, time_next):
-        batch, *_, device = *x.shape, x.device
-        return x
+    def sample(self, batch_size = 16):
+        shape = (batch_size, self.channels, self.image_size, self.image_size)
 
-    @torch.no_grad()
-    def sample_all_timesteps(self, shape):
         images = torch.randn(shape, device = self.device)
         steps = torch.linspace(1., 0., 100 + 1, device = self.device)
 
         for i in tqdm(range(100), desc = 'sampling loop time step', total = 100):
             times = steps[i]
             times_next = steps[i + 1]
-            images = self.sample_one_timestep(images, times, times_next)
+            images = images
 
         return unnormalize_to_zero_to_one(images)
 
-    @torch.no_grad()
-    def sample(self, batch_size = 16):
-        return self.sample_all_timesteps((batch_size, self.channels, self.image_size, self.image_size))
-
     # training
 
     def forward(self, images):

From 76b79aa8479ac8ab93626427eea3cfa5f3c57490 Mon Sep 17 00:00:00 2001
From: Phil Wang <lucidrains@gmail.com>
Date: Tue, 28 Jun 2022 14:16:05 -0700
Subject: [PATCH 05/10] take care of equation 7 in the paper

---
 .../elucidated_diffusion.py                   | 30 +++++++++++++++----
 1 file changed, 25 insertions(+), 5 deletions(-)

diff --git a/denoising_diffusion_pytorch/elucidated_diffusion.py b/denoising_diffusion_pytorch/elucidated_diffusion.py
index 1dfff27..c918950 100644
--- a/denoising_diffusion_pytorch/elucidated_diffusion.py
+++ b/denoising_diffusion_pytorch/elucidated_diffusion.py
@@ -37,6 +37,7 @@ class ElucidatedDiffusion(nn.Module):
         *,
         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
@@ -46,7 +47,7 @@ class ElucidatedDiffusion(nn.Module):
         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
+        S_noise = 1.003,
     ):
         super().__init__()
         assert net.learned_sinusoidal_cond
@@ -69,6 +70,8 @@ class ElucidatedDiffusion(nn.Module):
         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
@@ -101,6 +104,23 @@ class ElucidatedDiffusion(nn.Module):
     def loss_weight(self, sigma):
         return (sigma ** 2 + self.sigma_data ** 2) * (sigma * self.sigma_data) ** -2
 
+    # 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)
+
+        rho, sigma_max, sigma_min = self.rho, self.sigma_max, self.sigma_min
+
+        N = num_sample_steps
+        inv_rho = 1 / rho
+
+        for i in range(num_sample_steps - 1):
+            next_sigma = (sigma_max ** inv_rho + i / (N - 1) * (sigma_min ** inv_rho - sigma_max ** inv_rho)) ** rho
+            yield next_sigma
+
+        yield 0. # last step return 0.
+
     # preconditioned network output
     # equation (7) in the paper
 
@@ -121,11 +141,11 @@ class ElucidatedDiffusion(nn.Module):
         shape = (batch_size, self.channels, self.image_size, self.image_size)
 
         images = torch.randn(shape, device = self.device)
-        steps = torch.linspace(1., 0., 100 + 1, device = self.device)
 
-        for i in tqdm(range(100), desc = 'sampling loop time step', total = 100):
-            times = steps[i]
-            times_next = steps[i + 1]
+        sigma_schedule = [*self.sample_schedule()]
+        sigma_schedule = list(zip(sigma_schedule[:-1], sigma_schedule[1:]))
+
+        for sigma, sigma_next in tqdm(sigma_schedule, desc = 'sampling time step'):
             images = images
 
         return unnormalize_to_zero_to_one(images)

From 618493714f75570ac4c4c298abe94f9b7305079d Mon Sep 17 00:00:00 2001
From: Phil Wang <lucidrains@gmail.com>
Date: Tue, 28 Jun 2022 14:34:39 -0700
Subject: [PATCH 06/10] clamp the sigma coming out of the log normal
 distribution

---
 denoising_diffusion_pytorch/elucidated_diffusion.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/denoising_diffusion_pytorch/elucidated_diffusion.py b/denoising_diffusion_pytorch/elucidated_diffusion.py
index c918950..114dc3e 100644
--- a/denoising_diffusion_pytorch/elucidated_diffusion.py
+++ b/denoising_diffusion_pytorch/elucidated_diffusion.py
@@ -99,7 +99,8 @@ class ElucidatedDiffusion(nn.Module):
     # noise distribution
 
     def noise_distribution(self, batch_size):
-        return (self.P_mean + self.P_std * torch.randn((batch_size,), device = self.device)).exp()
+        sigmas = (self.P_mean + self.P_std * torch.randn((batch_size,), device = self.device)).exp()
+        return sigmas.clamp(min = self.sigma_min, max =self.sigma_max)
 
     def loss_weight(self, sigma):
         return (sigma ** 2 + self.sigma_data ** 2) * (sigma * self.sigma_data) ** -2

From f4b1d7a67c388626838e77dcb5ed30f683f485c9 Mon Sep 17 00:00:00 2001
From: Phil Wang <lucidrains@gmail.com>
Date: Tue, 28 Jun 2022 15:15:21 -0700
Subject: [PATCH 07/10] complete a first pass of elucidated ddpm

---
 .../elucidated_diffusion.py                   | 51 ++++++++++++++++---
 1 file changed, 43 insertions(+), 8 deletions(-)

diff --git a/denoising_diffusion_pytorch/elucidated_diffusion.py b/denoising_diffusion_pytorch/elucidated_diffusion.py
index 114dc3e..175d8be 100644
--- a/denoising_diffusion_pytorch/elucidated_diffusion.py
+++ b/denoising_diffusion_pytorch/elucidated_diffusion.py
@@ -1,3 +1,4 @@
+from math import sqrt
 import torch
 from torch import nn, einsum
 import torch.nn.functional as F
@@ -90,7 +91,7 @@ class ElucidatedDiffusion(nn.Module):
         return sigma * self.sigma_data * (self.sigma_data ** 2 + sigma ** 2) ** -0.5
 
     def c_in(self, sigma):
-        return (sigma ** 2 + self.sigma_data ** 2) * -0.5
+        return 1 * (sigma ** 2 + self.sigma_data ** 2) ** -0.5
 
     def c_noise(self, sigma):
         """ apparently empirically derived """
@@ -111,21 +112,28 @@ class ElucidatedDiffusion(nn.Module):
     def sample_schedule(self, num_sample_steps = None):
         num_sample_steps = default(num_sample_steps, self.num_sample_steps)
 
-        rho, sigma_max, sigma_min = self.rho, self.sigma_max, self.sigma_min
+        rho, sigma_max, sigma_min, S_tmin, S_tmax, S_churn = self.rho, self.sigma_max, self.sigma_min, self.S_tmin, self.S_tmax, self.S_churn
+        gamma = min(S_churn / num_sample_steps, sqrt(2) - 1)
 
         N = num_sample_steps
         inv_rho = 1 / rho
 
         for i in range(num_sample_steps - 1):
-            next_sigma = (sigma_max ** inv_rho + i / (N - 1) * (sigma_min ** inv_rho - sigma_max ** inv_rho)) ** rho
-            yield next_sigma
+            sigma_i = (sigma_max ** inv_rho + i / (N - 1) * (sigma_min ** inv_rho - sigma_max ** inv_rho)) ** rho
+            gamma_i = gamma if S_tmin <= sigma_i <= S_tmax else 0.
+            yield sigma_i, gamma_i
 
-        yield 0. # last step return 0.
+        yield 0., 0. # last step return 0.
 
     # preconditioned network output
     # equation (7) in the paper
 
     def preconditioned_network_forward(self, noised_images, sigma):
+        batch, device = noised_images.shape[0], noised_images.device
+
+        if isinstance(sigma, float):
+            sigma = torch.ones((batch,), device = device) * sigma
+
         padded_sigma = rearrange(sigma, 'b -> b 1 1 1')
 
         net_out = self.net(
@@ -141,13 +149,40 @@ class ElucidatedDiffusion(nn.Module):
     def sample(self, batch_size = 16):
         shape = (batch_size, self.channels, self.image_size, self.image_size)
 
-        images = torch.randn(shape, device = self.device)
+        # get the schedule, which is returned as (sigma, gamma) tuple, and pair up with the next sigma and gamma
 
         sigma_schedule = [*self.sample_schedule()]
         sigma_schedule = list(zip(sigma_schedule[:-1], sigma_schedule[1:]))
 
-        for sigma, sigma_next in tqdm(sigma_schedule, desc = 'sampling time step'):
-            images = images
+        # function to return noise, given a sigma value
+
+        get_noise = lambda std_dev: (std_dev * torch.randn(shape, device = self.device))
+
+        # images is None, set on first iteration
+
+        images = None
+
+        for (sigma, gamma), (sigma_next, gamma_next) in tqdm(sigma_schedule, desc = 'sampling time step'):
+            if not exists(images):
+                # images start off as the noise based off the first sigma
+                images = get_noise(sigma)
+
+            eps = get_noise(gamma)
+            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)
+            denoised = (images_hat - model_output) / sigma_hat
+
+            images_next = images_hat + (sigma_next - sigma_hat) * denoised
+
+            if sigma_next != 0:
+                # second order correction
+                model_output_next = self.preconditioned_network_forward(images_next, sigma_next)
+                denoised_prime = (images_next - model_output_next) / sigma_next
+                images_next = images_hat + (sigma_next - sigma_hat) * (0.5 * denoised + 0.5 * denoised_prime)
+
+            images = images_next
 
         return unnormalize_to_zero_to_one(images)
 

From a8403b83fe4bcfa9b7de5530850ee08f6152472d Mon Sep 17 00:00:00 2001
From: Phil Wang <lucidrains@gmail.com>
Date: Tue, 28 Jun 2022 16:00:16 -0700
Subject: [PATCH 08/10] no clamping when training from sigmas drawn from log
 normal distribution, clamp final images being sampled

---
 denoising_diffusion_pytorch/elucidated_diffusion.py | 9 ++++-----
 1 file changed, 4 insertions(+), 5 deletions(-)

diff --git a/denoising_diffusion_pytorch/elucidated_diffusion.py b/denoising_diffusion_pytorch/elucidated_diffusion.py
index 175d8be..782121c 100644
--- a/denoising_diffusion_pytorch/elucidated_diffusion.py
+++ b/denoising_diffusion_pytorch/elucidated_diffusion.py
@@ -94,14 +94,12 @@ class ElucidatedDiffusion(nn.Module):
         return 1 * (sigma ** 2 + self.sigma_data ** 2) ** -0.5
 
     def c_noise(self, sigma):
-        """ apparently empirically derived """
         return log(sigma) * 0.25
 
     # noise distribution
 
     def noise_distribution(self, batch_size):
-        sigmas = (self.P_mean + self.P_std * torch.randn((batch_size,), device = self.device)).exp()
-        return sigmas.clamp(min = self.sigma_min, max =self.sigma_max)
+        return (self.P_mean + self.P_std * torch.randn((batch_size,), device = self.device)).exp()
 
     def loss_weight(self, sigma):
         return (sigma ** 2 + self.sigma_data ** 2) * (sigma * self.sigma_data) ** -2
@@ -184,6 +182,7 @@ class ElucidatedDiffusion(nn.Module):
 
             images = images_next
 
+        images = images.clamp(-1., 1.)
         return unnormalize_to_zero_to_one(images)
 
     # training
@@ -203,9 +202,9 @@ class ElucidatedDiffusion(nn.Module):
 
         noised_images = images + padded_sigmas * noise  # alphas are 1. in the paper
 
-        model_out = self.preconditioned_network_forward(noised_images, sigmas)
+        denoised = self.preconditioned_network_forward(noised_images, sigmas)
 
-        losses = F.mse_loss(model_out, images, reduction = 'none')
+        losses = F.mse_loss(denoised, images, reduction = 'none')
         losses = reduce(losses, 'b ... -> b', 'mean')
 
         losses = losses * self.loss_weight(sigmas)

From b87ea2778111f40ba74fa2b7508c0d70a7628d2f Mon Sep 17 00:00:00 2001
From: Phil Wang <lucidrains@gmail.com>
Date: Tue, 28 Jun 2022 16:05:21 -0700
Subject: [PATCH 09/10] fix off by one

---
 denoising_diffusion_pytorch/elucidated_diffusion.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/denoising_diffusion_pytorch/elucidated_diffusion.py b/denoising_diffusion_pytorch/elucidated_diffusion.py
index 782121c..3204a97 100644
--- a/denoising_diffusion_pytorch/elucidated_diffusion.py
+++ b/denoising_diffusion_pytorch/elucidated_diffusion.py
@@ -116,7 +116,7 @@ class ElucidatedDiffusion(nn.Module):
         N = num_sample_steps
         inv_rho = 1 / rho
 
-        for i in range(num_sample_steps - 1):
+        for i in range(num_sample_steps):
             sigma_i = (sigma_max ** inv_rho + i / (N - 1) * (sigma_min ** inv_rho - sigma_max ** inv_rho)) ** rho
             gamma_i = gamma if S_tmin <= sigma_i <= S_tmax else 0.
             yield sigma_i, gamma_i

From 5db64fec4bc1c34dac5ecdde5d9e6ffdc7b32c97 Mon Sep 17 00:00:00 2001
From: Phil Wang <lucidrains@gmail.com>
Date: Tue, 28 Jun 2022 17:26:58 -0700
Subject: [PATCH 10/10] refactor sigmas and gamma generation

---
 .../elucidated_diffusion.py                   | 55 ++++++++++---------
 1 file changed, 29 insertions(+), 26 deletions(-)

diff --git a/denoising_diffusion_pytorch/elucidated_diffusion.py b/denoising_diffusion_pytorch/elucidated_diffusion.py
index 3204a97..2b3217e 100644
--- a/denoising_diffusion_pytorch/elucidated_diffusion.py
+++ b/denoising_diffusion_pytorch/elucidated_diffusion.py
@@ -110,18 +110,14 @@ class ElucidatedDiffusion(nn.Module):
     def sample_schedule(self, num_sample_steps = None):
         num_sample_steps = default(num_sample_steps, self.num_sample_steps)
 
-        rho, sigma_max, sigma_min, S_tmin, S_tmax, S_churn = self.rho, self.sigma_max, self.sigma_min, self.S_tmin, self.S_tmax, self.S_churn
-        gamma = min(S_churn / num_sample_steps, sqrt(2) - 1)
-
         N = num_sample_steps
-        inv_rho = 1 / rho
+        inv_rho = 1 / self.rho
 
-        for i in range(num_sample_steps):
-            sigma_i = (sigma_max ** inv_rho + i / (N - 1) * (sigma_min ** inv_rho - sigma_max ** inv_rho)) ** rho
-            gamma_i = gamma if S_tmin <= sigma_i <= S_tmax else 0.
-            yield sigma_i, gamma_i
+        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
 
-        yield 0., 0. # last step return 0.
+        sigmas = F.pad(sigmas, (0, 1), value = 0.) # last step is sigma value of 0.
+        return sigmas
 
     # preconditioned network output
     # equation (7) in the paper
@@ -130,7 +126,7 @@ class ElucidatedDiffusion(nn.Module):
         batch, device = noised_images.shape[0], noised_images.device
 
         if isinstance(sigma, float):
-            sigma = torch.ones((batch,), device = device) * sigma
+            sigma = torch.full((batch,), sigma, device = device)
 
         padded_sigma = rearrange(sigma, 'b -> b 1 1 1')
 
@@ -149,36 +145,43 @@ class ElucidatedDiffusion(nn.Module):
 
         # get the schedule, which is returned as (sigma, gamma) tuple, and pair up with the next sigma and gamma
 
-        sigma_schedule = [*self.sample_schedule()]
-        sigma_schedule = list(zip(sigma_schedule[:-1], sigma_schedule[1:]))
+        sigmas = self.sample_schedule()
 
-        # function to return noise, given a sigma value
+        gammas = torch.where(
+            (sigmas >= self.S_tmin) & (sigmas <= self.S_tmax),
+            min(self.S_churn / self.num_sample_steps, sqrt(2) - 1),
+            0.
+        )
 
-        get_noise = lambda std_dev: (std_dev * torch.randn(shape, device = self.device))
+        sigmas_and_gammas = list(zip(sigmas[:-1], sigmas[1:], gammas[:-1]))
 
-        # images is None, set on first iteration
+        # images is noise at the beginning
 
-        images = None
+        init_sigma = sigmas[0]
 
-        for (sigma, gamma), (sigma_next, gamma_next) in tqdm(sigma_schedule, desc = 'sampling time step'):
-            if not exists(images):
-                # images start off as the noise based off the first sigma
-                images = get_noise(sigma)
+        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 = gamma * torch.randn(shape, device = self.device)
 
-            eps = get_noise(gamma)
             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)
-            denoised = (images_hat - model_output) / sigma_hat
+            denoised_over_sigma = (images_hat - model_output) / sigma_hat
 
-            images_next = images_hat + (sigma_next - sigma_hat) * denoised
+            images_next = images_hat + (sigma_next - sigma_hat) * denoised_over_sigma
+
+            # second order correction, if not the last timestep
 
             if sigma_next != 0:
-                # second order correction
                 model_output_next = self.preconditioned_network_forward(images_next, sigma_next)
-                denoised_prime = (images_next - model_output_next) / sigma_next
-                images_next = images_hat + (sigma_next - sigma_hat) * (0.5 * denoised + 0.5 * denoised_prime)
+                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