SimpleDiffusionModels/samplers.py at master · okumoto-sho/SimpleDiffusionModels · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
import torch

class DDPMSampler1d():
    """
        Denoising Diffusion Probabilistic Model (DDPM) sampler.
        Reference: https://arxiv.org/pdf/2006.11239.pdf

        Parameters
        ----------
        betas: torch.Tensor of shape [T]
            Beta values for each time step
    """

    def __init__(self, betas: torch.Tensor, noise_model):
        self.device = betas.device
        self.betas = betas
        self.alphas = 1 - betas
        self.alpha_bars = torch.cumprod(self.alphas, dim=0).to(self.device)
        self.noise_model = noise_model

    def sample(self, x, t):
        """
            Sample from the DDPM model.

            Parameters
            ----------
            x: torch.tensor of shape [N, C, D]
                Sample of previous time step

            t: torch.tensor of shape [N]
                Time step

            Returns
            -------
            samples: torch.Tensor
                Sample at t - 1 from the DDPM model
        """

        scale = (1 - self.alphas[t+1]) / torch.sqrt(1 - self.alpha_bars[t+1])
        sigma =  (1 - self.alpha_bars[t]) * (1 - self.alphas[t+1]) / (1 - self.alpha_bars[t+1])
        inv = (1 / torch.sqrt(self.alphas[t+1]))

        scale = scale[:, None, None]
        sigma = sigma[:, None, None]
        inv = inv[:, None, None]

        sample = inv * (x - scale * self.noise_model(x, t+1)) + torch.sqrt(sigma) * torch.randn_like(x).to(self.device)
        return sample

class DDPMSampler2d():
    """
        Denoising Diffusion Probabilistic Model (DDPM) sampler for 2D data.
        Reference: https://arxiv.org/pdf/2006.11239.pdf

        Parameters
        ----------
        betas: torch.Tensor of shape [T]
            Beta values for each time step
    """

    def __init__(self, betas: torch.Tensor, noise_model):
        self.device = betas.device
        self.betas = betas
        self.alphas = 1 - betas
        self.alpha_bars = torch.cumprod(self.alphas, dim=0).to(self.device)
        self.noise_model = noise_model

    def sample(self, x, t):
        """
            Sample from the DDPM model for 2D data.

            Parameters
            ----------
            x: torch.tensor of shape [N, C, H, W]
                2D Sample of previous time step

            t: torch.tensor of shape [N]
                Time step

            Returns
            -------
            samples: torch.Tensor of shape [N, C, H, W]
                2D Sample at t - 1 from the DDPM model
        """

        scale = (1 - self.alphas[t+1]) / torch.sqrt(1 - self.alpha_bars[t+1])
        sigma =  (1 - self.alpha_bars[t]) * (1 - self.alphas[t+1]) / (1 - self.alpha_bars[t+1])
        inv = (1 / torch.sqrt(self.alphas[t+1]))

        scale = scale[:, None, None, None]
        sigma = sigma[:, None, None, None]
        inv = inv[:, None, None, None]

        sample = inv * (x - scale * self.noise_model(x, t+1)) + torch.sqrt(sigma) * torch.randn_like(x).to(self.device)
        return sample