scheduling_ddim.py 6.59 KB
Newer Older
1
# Copyright 2022 Stanford University Team and The HuggingFace Team. All rights reserved.
Patrick von Platen's avatar
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
17

# DISCLAIMER: This code is strongly influenced by https://github.com/pesser/pytorch_diffusion
# and https://github.com/hojonathanho/diffusion

Patrick von Platen's avatar
Patrick von Platen committed
18
19
import math

Patrick von Platen's avatar
Patrick von Platen committed
20
import numpy as np
Patrick von Platen's avatar
Patrick von Platen committed
21
22

from ..configuration_utils import ConfigMixin
23
24
25
26
27
from .scheduling_utils import SchedulerMixin


def betas_for_alpha_bar(num_diffusion_timesteps, max_beta=0.999):
    """
Patrick von Platen's avatar
Patrick von Platen committed
28
29
    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].
30

Patrick von Platen's avatar
Patrick von Platen committed
31
32
33
    :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.
34
35
36
    :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
Patrick von Platen committed
47
48


Patrick von Platen's avatar
Patrick von Platen committed
49
class DDIMScheduler(SchedulerMixin, ConfigMixin):
Patrick von Platen's avatar
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",
patil-suraj's avatar
patil-suraj committed
56
57
        trained_betas=None,
        timestep_values=None,
Patrick von Platen's avatar
Patrick von Platen committed
58
        clip_sample=True,
Patrick von Platen's avatar
Patrick von Platen committed
59
        tensor_format="np",
Patrick von Platen's avatar
Patrick von Platen committed
60
61
    ):
        super().__init__()
62
        self.register_to_config(
Patrick von Platen's avatar
Patrick von Platen committed
63
64
65
66
            timesteps=timesteps,
            beta_start=beta_start,
            beta_end=beta_end,
            beta_schedule=beta_schedule,
67
68
69
            trained_betas=trained_betas,
            timestep_values=timestep_values,
            clip_sample=clip_sample,
Patrick von Platen's avatar
Patrick von Platen committed
70
71
        )

72
        if beta_schedule == "linear":
73
            self.betas = np.linspace(beta_start, beta_end, timesteps, dtype=np.float32)
74
75
76
        elif beta_schedule == "scaled_linear":
            # this schedule is very specific to the latent diffusion model.
            self.betas = np.linspace(beta_start**0.5, beta_end**0.5, timesteps, dtype=np.float32) ** 2
Patrick von Platen's avatar
Patrick von Platen committed
77
        elif beta_schedule == "squaredcos_cap_v2":
Patrick von Platen's avatar
Patrick von Platen committed
78
            # Glide cosine schedule
79
            self.betas = betas_for_alpha_bar(timesteps)
Patrick von Platen's avatar
Patrick von Platen committed
80
81
82
        else:
            raise NotImplementedError(f"{beta_schedule} does is not implemented for {self.__class__}")

Patrick von Platen's avatar
Patrick von Platen committed
83
84
85
86
87
88
        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)

Patrick von Platen's avatar
Patrick von Platen committed
89
    def get_variance(self, t, num_inference_steps):
90
91
        orig_t = self.config.timesteps // num_inference_steps * t
        orig_prev_t = self.config.timesteps // num_inference_steps * (t - 1) if t > 0 else -1
Patrick von Platen's avatar
Patrick von Platen committed
92

93
94
        alpha_prod_t = self.alphas_cumprod[orig_t]
        alpha_prod_t_prev = self.alphas_cumprod[orig_prev_t] if orig_prev_t >= 0 else self.one
Patrick von Platen's avatar
Patrick von Platen committed
95
96
97
98
99
100
101
        beta_prod_t = 1 - alpha_prod_t
        beta_prod_t_prev = 1 - alpha_prod_t_prev

        variance = (beta_prod_t_prev / beta_prod_t) * (1 - alpha_prod_t / alpha_prod_t_prev)

        return variance

102
    def step(self, residual, sample, t, num_inference_steps, eta, use_clipped_residual=False):
Patrick von Platen's avatar
Patrick von Platen committed
103
104
105
106
107
        # See formulas (12) and (16) of DDIM paper https://arxiv.org/pdf/2010.02502.pdf
        # Ideally, read DDIM paper in-detail understanding

        # Notation (<variable name> -> <name in paper>
        # - pred_noise_t -> e_theta(x_t, t)
108
        # - pred_original_sample -> f_theta(x_t, t) or x_0
Patrick von Platen's avatar
Patrick von Platen committed
109
110
        # - std_dev_t -> sigma_t
        # - eta -> η
111
112
        # - pred_sample_direction -> "direction pointingc to x_t"
        # - pred_prev_sample -> "x_t-1"
Patrick von Platen's avatar
Patrick von Platen committed
113
114

        # 1. get actual t and t-1
115
116
        orig_t = self.config.timesteps // num_inference_steps * t
        orig_prev_t = self.config.timesteps // num_inference_steps * (t - 1) if t > 0 else -1
Patrick von Platen's avatar
Patrick von Platen committed
117
118

        # 2. compute alphas, betas
119
120
        alpha_prod_t = self.alphas_cumprod[orig_t]
        alpha_prod_t_prev = self.alphas_cumprod[orig_prev_t] if orig_prev_t >= 0 else self.one
Patrick von Platen's avatar
Patrick von Platen committed
121
122
        beta_prod_t = 1 - alpha_prod_t

123
        # 3. compute predicted original sample from predicted noise also called
Patrick von Platen's avatar
Patrick von Platen committed
124
        # "predicted x_0" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf
125
        pred_original_sample = (sample - beta_prod_t ** (0.5) * residual) / alpha_prod_t ** (0.5)
Patrick von Platen's avatar
Patrick von Platen committed
126
127

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

        # 5. compute variance: "sigma_t(η)" -> see formula (16)
        # σ_t = sqrt((1 − α_t−1)/(1 − α_t)) * sqrt(1 − α_t/α_t−1)
        variance = self.get_variance(t, num_inference_steps)
Patrick von Platen's avatar
Patrick von Platen committed
134
        std_dev_t = eta * variance ** (0.5)
Patrick von Platen's avatar
Patrick von Platen committed
135

anton-l's avatar
anton-l committed
136
        if use_clipped_residual:
Patrick von Platen's avatar
Patrick von Platen committed
137
            # the residual is always re-derived from the clipped x_0 in Glide
138
            residual = (sample - alpha_prod_t ** (0.5) * pred_original_sample) / beta_prod_t ** (0.5)
anton-l's avatar
anton-l committed
139

Patrick von Platen's avatar
Patrick von Platen committed
140
        # 6. compute "direction pointing to x_t" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf
141
        pred_sample_direction = (1 - alpha_prod_t_prev - std_dev_t**2) ** (0.5) * residual
Patrick von Platen's avatar
Patrick von Platen committed
142
143

        # 7. compute x_t without "random noise" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf
144
        pred_prev_sample = alpha_prod_t_prev ** (0.5) * pred_original_sample + pred_sample_direction
Patrick von Platen's avatar
Patrick von Platen committed
145

146
        return pred_prev_sample
Patrick von Platen's avatar
Patrick von Platen committed
147

148
149
150
151
152
153
154
155
156
    def add_noise(self, original_samples, noise, timesteps):
        sqrt_alpha_prod = self.alphas_cumprod[timesteps] ** 0.5
        sqrt_alpha_prod = self.match_shape(sqrt_alpha_prod, original_samples)
        sqrt_one_minus_alpha_prod = (1 - self.alphas_cumprod[timesteps]) ** 0.5
        sqrt_one_minus_alpha_prod = self.match_shape(sqrt_one_minus_alpha_prod, original_samples)

        noisy_samples = sqrt_alpha_prod * original_samples + sqrt_one_minus_alpha_prod * noise
        return noisy_samples

Patrick von Platen's avatar
Patrick von Platen committed
157
    def __len__(self):
158
        return self.config.timesteps