scheduling_ddpm.py 6.02 KB
Newer Older
1
# Copyright 2022 UC Berkely Team and The HuggingFace Team. All rights reserved.
Patrick von Platen's avatar
improve  
Patrick von Platen committed
2
3
4
5
6
7
8
9
10
11
12
13
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
14
15
16

# DISCLAIMER: This file is strongly influenced by https://github.com/ermongroup/ddim

anton-l's avatar
anton-l committed
17
import math
Patrick von Platen's avatar
Patrick von Platen committed
18

Patrick von Platen's avatar
Patrick von Platen committed
19
import numpy as np
Patrick von Platen's avatar
improve  
Patrick von Platen committed
20

Patrick von Platen's avatar
Patrick von Platen committed
21
from ..configuration_utils import ConfigMixin
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
from .scheduling_utils import SchedulerMixin


def betas_for_alpha_bar(num_diffusion_timesteps, max_beta=0.999):
    """
    Create a beta schedule that discretizes the given alpha_t_bar function,
    which defines the cumulative product of (1-beta) over time from t = [0,1].

    :param num_diffusion_timesteps: the number of betas to produce.
    :param alpha_bar: a lambda that takes an argument t from 0 to 1 and
                      produces the cumulative product of (1-beta) up to that
                      part of the diffusion process.
    :param max_beta: the maximum beta to use; use values lower than 1 to
                     prevent singularities.
    """
37

38
39
40
41
42
43
44
45
46
    def alpha_bar(time_step):
        return math.cos((time_step + 0.008) / 1.008 * math.pi / 2) ** 2

    betas = []
    for i in range(num_diffusion_timesteps):
        t1 = i / num_diffusion_timesteps
        t2 = (i + 1) / num_diffusion_timesteps
        betas.append(min(1 - alpha_bar(t2) / alpha_bar(t1), max_beta))
    return np.array(betas, dtype=np.float32)
Patrick von Platen's avatar
improve  
Patrick von Platen committed
47
48


Patrick von Platen's avatar
Patrick von Platen committed
49
class DDPMScheduler(SchedulerMixin, ConfigMixin):
Patrick von Platen's avatar
improve  
Patrick von Platen committed
50
51
52
53
54
55
    def __init__(
        self,
        timesteps=1000,
        beta_start=0.0001,
        beta_end=0.02,
        beta_schedule="linear",
56
57
        trained_betas=None,
        timestep_values=None,
Patrick von Platen's avatar
improve  
Patrick von Platen committed
58
        variance_type="fixed_small",
Patrick von Platen's avatar
Patrick von Platen committed
59
        clip_sample=True,
Patrick von Platen's avatar
Patrick von Platen committed
60
        tensor_format="np",
Patrick von Platen's avatar
improve  
Patrick von Platen committed
61
62
    ):
        super().__init__()
63
        self.register_to_config(
Patrick von Platen's avatar
improve  
Patrick von Platen committed
64
65
66
67
            timesteps=timesteps,
            beta_start=beta_start,
            beta_end=beta_end,
            beta_schedule=beta_schedule,
68
69
            trained_betas=trained_betas,
            timestep_values=timestep_values,
Patrick von Platen's avatar
improve  
Patrick von Platen committed
70
            variance_type=variance_type,
Patrick von Platen's avatar
Patrick von Platen committed
71
            clip_sample=clip_sample,
Patrick von Platen's avatar
improve  
Patrick von Platen committed
72
73
        )

74
75
76
        if trained_betas is not None:
            self.betas = np.asarray(trained_betas)
        elif beta_schedule == "linear":
77
            self.betas = np.linspace(beta_start, beta_end, timesteps, dtype=np.float32)
anton-l's avatar
anton-l committed
78
79
        elif beta_schedule == "squaredcos_cap_v2":
            # GLIDE cosine schedule
80
            self.betas = betas_for_alpha_bar(timesteps)
Patrick von Platen's avatar
improve  
Patrick von Platen committed
81
82
83
        else:
            raise NotImplementedError(f"{beta_schedule} does is not implemented for {self.__class__}")

Patrick von Platen's avatar
Patrick von Platen committed
84
85
86
87
88
89
        self.alphas = 1.0 - self.betas
        self.alphas_cumprod = np.cumprod(self.alphas, axis=0)
        self.one = np.array(1.0)

        self.set_format(tensor_format=tensor_format)

90
    def get_variance(self, t, variance_type=None):
91
92
        alpha_prod_t = self.alphas_cumprod[t]
        alpha_prod_t_prev = self.alphas_cumprod[t - 1] if t > 0 else self.one
Patrick von Platen's avatar
Patrick von Platen committed
93
94

        # For t > 0, compute predicted variance βt (see formala (6) and (7) from https://arxiv.org/pdf/2006.11239.pdf)
95
96
        # and sample from it to get previous sample
        # x_{t-1} ~ N(pred_prev_sample, variance) == add variane to pred_sample
97
        variance = (1 - alpha_prod_t_prev) / (1 - alpha_prod_t) * self.betas[t]
Patrick von Platen's avatar
Patrick von Platen committed
98

99
100
101
        if variance_type is None:
            variance_type = self.config.variance_type

Patrick von Platen's avatar
Patrick von Platen committed
102
        # hacks - were probs added for training stability
103
        if variance_type == "fixed_small":
Patrick von Platen's avatar
Patrick von Platen committed
104
            variance = self.clip(variance, min_value=1e-20)
105
        # for rl-diffuser https://arxiv.org/abs/2205.09991
106
        elif variance_type == "fixed_small_log":
107
            variance = self.log(self.clip(variance, min_value=1e-20))
108
        elif variance_type == "fixed_large":
109
            variance = self.betas[t]
110
111
112
        elif variance_type == "fixed_large_log":
            # GLIDE max_log
            variance = self.log(self.betas[t])
Patrick von Platen's avatar
Patrick von Platen committed
113
114
115

        return variance

116
    def step(self, residual, sample, t, predict_epsilon=True):
Patrick von Platen's avatar
Patrick von Platen committed
117
        # 1. compute alphas, betas
118
119
        alpha_prod_t = self.alphas_cumprod[t]
        alpha_prod_t_prev = self.alphas_cumprod[t - 1] if t > 0 else self.one
Patrick von Platen's avatar
Patrick von Platen committed
120
121
122
        beta_prod_t = 1 - alpha_prod_t
        beta_prod_t_prev = 1 - alpha_prod_t_prev

123
        # 2. compute predicted original sample from predicted noise also called
Patrick von Platen's avatar
Patrick von Platen committed
124
        # "predicted x_0" of formula (15) from https://arxiv.org/pdf/2006.11239.pdf
125
126
127
128
        if predict_epsilon:
            pred_original_sample = (sample - beta_prod_t ** (0.5) * residual) / alpha_prod_t ** (0.5)
        else:
            pred_original_sample = residual
Patrick von Platen's avatar
Patrick von Platen committed
129
130

        # 3. Clip "predicted x_0"
131
        if self.config.clip_sample:
132
            pred_original_sample = self.clip(pred_original_sample, -1, 1)
Patrick von Platen's avatar
Patrick von Platen committed
133

134
        # 4. Compute coefficients for pred_original_sample x_0 and current sample x_t
Patrick von Platen's avatar
Patrick von Platen committed
135
        # See formula (7) from https://arxiv.org/pdf/2006.11239.pdf
136
137
        pred_original_sample_coeff = (alpha_prod_t_prev ** (0.5) * self.betas[t]) / beta_prod_t
        current_sample_coeff = self.alphas[t] ** (0.5) * beta_prod_t_prev / beta_prod_t
Patrick von Platen's avatar
Patrick von Platen committed
138

139
        # 5. Compute predicted previous sample µ_t
Patrick von Platen's avatar
Patrick von Platen committed
140
        # See formula (7) from https://arxiv.org/pdf/2006.11239.pdf
141
        pred_prev_sample = pred_original_sample_coeff * pred_original_sample + current_sample_coeff * sample
Patrick von Platen's avatar
Patrick von Platen committed
142

143
        return pred_prev_sample
Patrick von Platen's avatar
Patrick von Platen committed
144

145
    def forward_step(self, original_sample, noise, t):
anton-l's avatar
anton-l committed
146
147
        sqrt_alpha_prod = self.alphas_cumprod[t] ** 0.5
        sqrt_one_minus_alpha_prod = (1 - self.alphas_cumprod[t]) ** 0.5
148
149
        noisy_sample = sqrt_alpha_prod * original_sample + sqrt_one_minus_alpha_prod * noise
        return noisy_sample
anton-l's avatar
anton-l committed
150

Patrick von Platen's avatar
improve  
Patrick von Platen committed
151
    def __len__(self):
152
        return self.config.timesteps