#

Denoising Diffusion Probabilistic Models (DDPM) evaluation/sampling

This is the code to generate images and create interpolations between given images.

14import numpy as np
15import torch
16from matplotlib import pyplot as plt
17from torchvision.transforms.functional import to_pil_image, resize
18
19from labml import experiment, monit
20from labml_nn.diffusion.ddpm import DenoiseDiffusion, gather
21from labml_nn.diffusion.ddpm.experiment import Configs

#

Sampler class

24class Sampler:

#

diffusion is the DenoiseDiffusion instance
image_channels is the number of channels in the image
image_size is the image size
device is the device of the model

29    def __init__(self, diffusion: DenoiseDiffusion, image_channels: int, image_size: int, device: torch.device):

#

36        self.device = device
37        self.image_size = image_size
38        self.image_channels = image_channels
39        self.diffusion = diffusion

#

$T$

42        self.n_steps = diffusion.n_steps

#

$ϵ_{θ} (x_{t}, t)$

44        self.eps_model = diffusion.eps_model

#

$β_{t}$

46        self.beta = diffusion.beta

#

$α_{t}$

48        self.alpha = diffusion.alpha

#

$\overset{α_{t}}{ˉ}$

50        self.alpha_bar = diffusion.alpha_bar

#

$\overset{α}{ˉ}_{t - 1}$

52        alpha_bar_tm1 = torch.cat([self.alpha_bar.new_ones((1,)), self.alpha_bar[:-1]])

#

To calculate

q (x_{t - 1} ∣ x_{t}, x_{0}) \tilde{μ}_{t} (x_{t}, x_{0}) \tilde{β_{t}} = N (x_{t - 1}; \tilde{μ}_{t} (x_{t}, x_{0}), \tilde{β_{t}} I) = \frac{α ˉ _{t - 1} β _{t}}{1 - α _{t} ˉ} x_{0} + \frac{α _{t} ( 1 - α ˉ _{t - 1} )}{1 - α _{t} ˉ} x_{t} = \frac{1 - α ˉ _{t - 1}}{1 - α _{t} ˉ} β_{t}

#

$\tilde{β_{t}} = \frac{1 - α ˉ _{t - 1}}{1 - α _{t} ˉ} β_{t}$

64        self.beta_tilde = self.beta * (1 - alpha_bar_tm1) / (1 - self.alpha_bar)

#

$\frac{α ˉ _{t - 1} β _{t}}{1 - α _{t} ˉ}$

66        self.mu_tilde_coef1 = self.beta * (alpha_bar_tm1 ** 0.5) / (1 - self.alpha_bar)

#

$\frac{α _{t} ( 1 - α ˉ _{t - 1}}{1 - α _{t} ˉ}$

68        self.mu_tilde_coef2 = (self.alpha ** 0.5) * (1 - alpha_bar_tm1) / (1 - self.alpha_bar)

#

$σ^{2} = β$

70        self.sigma2 = self.beta

#

Helper function to display an image

72    def show_image(self, img, title=""):

#

74        img = img.clip(0, 1)
75        img = img.cpu().numpy()
76        plt.imshow(img.transpose(1, 2, 0))
77        plt.title(title)
78        plt.show()

#

Helper function to create a video

80    def make_video(self, frames, path="video.mp4"):

#

82        import imageio

#

20 second video

84        writer = imageio.get_writer(path, fps=len(frames) // 20)

#

Add each image

86        for f in frames:
87            f = f.clip(0, 1)
88            f = to_pil_image(resize(f, [368, 368]))
89            writer.append_data(np.array(f))

#

91        writer.close()

#

Sample an image step-by-step using $p_{θ} (x_{t - 1} ∣ x_{t})$

We sample an image step-by-step using $p_{θ} (x_{t - 1} ∣ x_{t})$ and at each step show the estimate $x_{0} \approx \overset{x}{^}_{0} = \frac{1}{α ˉ} (x_{t} - 1 - \overset{α_{t}}{ˉ} ϵ_{θ} (x_{t}, t))$

93    def sample_animation(self, n_frames: int = 1000, create_video: bool = True):

#

$x_{T} \sim p (x_{T}) = N (x_{T}; 0, I)$

104        xt = torch.randn([1, self.image_channels, self.image_size, self.image_size], device=self.device)

#

Interval to log $\overset{x}{^}_{0}$

107        interval = self.n_steps // n_frames

#

Frames for video

109        frames = []

#

Sample $T$ steps

111        for t_inv in monit.iterate('Denoise', self.n_steps):

#

$t$

113            t_ = self.n_steps - t_inv - 1

#

$t$ in a tensor

115            t = xt.new_full((1,), t_, dtype=torch.long)

#

$ϵ_{θ} (x_{t}, t)$

117            eps_theta = self.eps_model(xt, t)
118            if t_ % interval == 0:

#

Get $\overset{x}{^}_{0}$ and add to frames

120                x0 = self.p_x0(xt, t, eps_theta)
121                frames.append(x0[0])
122                if not create_video:
123                    self.show_image(x0[0], f"{t_}")

#

Sample from $p_{θ} (x_{t - 1} ∣ x_{t})$

125            xt = self.p_sample(xt, t, eps_theta)

#

Make video

128        if create_video:
129            self.make_video(frames)

#

Interpolate two images $x_{0}$ and $x_{0}^{'}$

We get $x_{t} \sim q (x_{t} ∣ x_{0})$ and $x_{t}^{'} \sim q (x_{t}^{'} ∣ x_{0})$ .

Then interpolate to $\overset{x}{ˉ}_{t} = (1 - λ) x_{t} + λ x_{0}^{'}$

Then get $\overset{x}{ˉ}_{0} \sim p_{θ} (x_{0} ∣ \overset{x}{ˉ}_{t})$

x1 is $x_{0}$
x2 is $x_{0}^{'}$
lambda_ is $λ$
t_ is $t$

131    def interpolate(self, x1: torch.Tensor, x2: torch.Tensor, lambda_: float, t_: int = 100):

#

Number of samples

150        n_samples = x1.shape[0]

#

$t$ tensor

152        t = torch.full((n_samples,), t_, device=self.device)

#

$\overset{x}{ˉ}_{t} = (1 - λ) x_{t} + λ x_{0}^{'}$

154        xt = (1 - lambda_) * self.diffusion.q_sample(x1, t) + lambda_ * self.diffusion.q_sample(x2, t)

#

$\overset{x}{ˉ}_{0} \sim p_{θ} (x_{0} ∣ \overset{x}{ˉ}_{t})$

157        return self._sample_x0(xt, t_)

#

Interpolate two images $x_{0}$ and $x_{0}^{'}$ and make a video

x1 is $x_{0}$
x2 is $x_{0}^{'}$
n_frames is the number of frames for the image
t_ is $t$
create_video specifies whether to make a video or to show each frame

159    def interpolate_animate(self, x1: torch.Tensor, x2: torch.Tensor, n_frames: int = 100, t_: int = 100,
160                            create_video=True):

#

Show original images

172        self.show_image(x1, "x1")
173        self.show_image(x2, "x2")

#

Add batch dimension

175        x1 = x1[None, :, :, :]
176        x2 = x2[None, :, :, :]

#

$t$ tensor

178        t = torch.full((1,), t_, device=self.device)

#

$x_{t} \sim q (x_{t} ∣ x_{0})$

180        x1t = self.diffusion.q_sample(x1, t)

#

$x_{t}^{'} \sim q (x_{t}^{'} ∣ x_{0})$

182        x2t = self.diffusion.q_sample(x2, t)
183
184        frames = []

#

Get frames with different $λ$

186        for i in monit.iterate('Interpolate', n_frames + 1, is_children_silent=True):

#

$λ$

188            lambda_ = i / n_frames

#

$\overset{x}{ˉ}_{t} = (1 - λ) x_{t} + λ x_{0}^{'}$

190            xt = (1 - lambda_) * x1t + lambda_ * x2t

#

$\overset{x}{ˉ}_{0} \sim p_{θ} (x_{0} ∣ \overset{x}{ˉ}_{t})$

192            x0 = self._sample_x0(xt, t_)

#

Add to frames

194            frames.append(x0[0])

#

Show frame

196            if not create_video:
197                self.show_image(x0[0], f"{lambda_ :.2f}")

#

Make video

200        if create_video:
201            self.make_video(frames)

#

Sample an image using $p_{θ} (x_{t - 1} ∣ x_{t})$

xt is $x_{t}$
n_steps is $t$

203    def _sample_x0(self, xt: torch.Tensor, n_steps: int):

#

Number of sampels

212        n_samples = xt.shape[0]

#

Iterate until $t$ steps

214        for t_ in monit.iterate('Denoise', n_steps):
215            t = n_steps - t_ - 1

#

Sample from $p_{θ} (x_{t - 1} ∣ x_{t})$

217            xt = self.diffusion.p_sample(xt, xt.new_full((n_samples,), t, dtype=torch.long))

#

Return $x_{0}$

220        return xt

#

Generate images

222    def sample(self, n_samples: int = 16):

#

$x_{T} \sim p (x_{T}) = N (x_{T}; 0, I)$

227        xt = torch.randn([n_samples, self.image_channels, self.image_size, self.image_size], device=self.device)

#

$x_{0} \sim p_{θ} (x_{0} ∣ x_{t})$

230        x0 = self._sample_x0(xt, self.n_steps)

#

Show images

233        for i in range(n_samples):
234            self.show_image(x0[i])

#

Sample from $p_{θ} (x_{t - 1} ∣ x_{t})$

p_{θ} (x_{t - 1} ∣ x_{t}) μ_{θ} (x_{t}, t) = N (x_{t - 1}; μ_{θ} (x_{t}, t), σ_{t}^{2} I) = \frac{1}{α _{t}} (x_{t} - \frac{β _{t}}{1 - α _{t} ˉ} ϵ_{θ} (x_{t}, t))

236    def p_sample(self, xt: torch.Tensor, t: torch.Tensor, eps_theta: torch.Tensor):

#

gather $\overset{α_{t}}{ˉ}$

249        alpha_bar = gather(self.alpha_bar, t)

#

$α_{t}$

251        alpha = gather(self.alpha, t)

#

$\frac{β}{1 - α _{t} ˉ}$

253        eps_coef = (1 - alpha) / (1 - alpha_bar) ** .5

#

$\frac{1}{α _{t}} (x_{t} - \frac{β _{t}}{1 - α _{t} ˉ} ϵ_{θ} (x_{t}, t))$

256        mean = 1 / (alpha ** 0.5) * (xt - eps_coef * eps_theta)

#

$σ^{2}$

258        var = gather(self.sigma2, t)

#

$ϵ \sim N (0, I)$

261        eps = torch.randn(xt.shape, device=xt.device)

#

Sample

263        return mean + (var ** .5) * eps

#

Estimate $x_{0}$

$x_{0} \approx \overset{x}{^}_{0} = \frac{1}{α ˉ} (x_{t} - 1 - \overset{α_{t}}{ˉ} ϵ_{θ} (x_{t}, t))$

265    def p_x0(self, xt: torch.Tensor, t: torch.Tensor, eps: torch.Tensor):

#

gather $\overset{α_{t}}{ˉ}$

273        alpha_bar = gather(self.alpha_bar, t)

#

$x_{0} \approx \overset{x}{^}_{0} = \frac{1}{α ˉ} (x_{t} - 1 - \overset{α_{t}}{ˉ} ϵ_{θ} (x_{t}, t))$

277        return (xt - (1 - alpha_bar) ** 0.5 * eps) / (alpha_bar ** 0.5)

#

Generate samples

280def main():

#

Training experiment run UUID

284    run_uuid = "a44333ea251411ec8007d1a1762ed686"

#

Start an evaluation

287    experiment.evaluate()

#

Create configs

290    configs = Configs()

#

Load custom configuration of the training run

292    configs_dict = experiment.load_configs(run_uuid)

#

Set configurations

294    experiment.configs(configs, configs_dict)

#

Initialize

297    configs.init()

#

Set PyTorch modules for saving and loading

300    experiment.add_pytorch_models({'eps_model': configs.eps_model})

#

Load training experiment

303    experiment.load(run_uuid)

#

Create sampler

306    sampler = Sampler(diffusion=configs.diffusion,
307                      image_channels=configs.image_channels,
308                      image_size=configs.image_size,
309                      device=configs.device)

#

Start evaluation

312    with experiment.start():

#

No gradients

314        with torch.no_grad():

#

Sample an image with an denoising animation

316            sampler.sample_animation()
317
318            if False:

#

Get some images fro data

320                data = next(iter(configs.data_loader)).to(configs.device)

#

Create an interpolation animation

323                sampler.interpolate_animate(data[0], data[1])

#

327if __name__ == '__main__':
328    main()

Denoising Diffusion Probabilistic Models (DDPM) evaluation/sampling

Sampler class

Sample an image step-by-step using pθ​(xt−1​∣xt​)

Interpolate two images x0​ and x0′​

Interpolate two images x0​ and x0′​ and make a video

Sample an image using pθ​(xt−1​∣xt​)

Generate images

Sample from pθ​(xt−1​∣xt​)

Estimate x0​

Sample an image step-by-step using $p_{θ} (x_{t - 1} ∣ x_{t})$

Interpolate two images $x_{0}$ and $x_{0}^{'}$

Interpolate two images $x_{0}$ and $x_{0}^{'}$ and make a video

Sample an image using $p_{θ} (x_{t - 1} ∣ x_{t})$

Sample from $p_{θ} (x_{t - 1} ∣ x_{t})$

Estimate $x_{0}$