SimpleDiffusionModels/diffusers.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
import torch

class DDPMDiffuser1d():
    """
        Denoising Diffusion Probabilistic Model (DDPM) diffuser.
        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):
        self.device = betas.device
        self.betas = betas
        self.alphas = 1 - betas
        self.alpha_bars = torch.cumprod(self.alphas, dim=0).to(self.device)

    def diffuse(self, X_0, t):
        """
            Diffuse data X_0 at time step t.

            Parameters
            ----------
            t: torch.Tensor of shape [N]
                Time step
            X_0: torch.Tensor of shape [N, C, D]
                Data to be diffused

            Returns
            -------
            X_t: torch.Tensor of shape [N, C, D]
                Diffused data
            eps: torch.Tensor of shape [N, C, D]
                Noise added to data
        """

        eps = torch.randn(X_0.size(), dtype=torch.float32, device=self.device)

        mean_scale = torch.sqrt(self.alpha_bars[t])[:, None, None]
        var_scale = torch.sqrt(1 - self.alpha_bars[t])[:, None, None]
        X_t = mean_scale * X_0 + var_scale * eps
        return X_t, eps

class DDPMDiffuser2d():
    """
        Denoising Diffusion Probabilistic Model (DDPM) diffuser 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):
        self.device = betas.device
        self.betas = betas
        self.alphas = 1 - betas
        self.alpha_bars = torch.cumprod(self.alphas, dim=0).to(self.device)

    def diffuse(self, X_0, t):
        """
            Diffuse 2D data X_0 at time step t.

            Parameters
            ----------
            t: torch.Tensor of shape [N]
                Time step
            X_0: torch.Tensor of shape [N, C, H, W]
                2D Data to be diffused

            Returns
            -------
            X_t: torch.Tensor of shape [N, C, H, W]
                Diffused 2D data
            eps: torch.Tensor of shape [N, C, H, W]
                Noise added to 2D data
        """

        eps = torch.randn(X_0.size(), dtype=torch.float32, device=self.device)

        mean_scale = torch.sqrt(self.alpha_bars[t])[:, None, None, None]
        var_scale = torch.sqrt(1 - self.alpha_bars[t])[:, None, None, None]
        X_t = mean_scale * X_0 + var_scale * eps
        return X_t, eps