scheduling_ddim.py 6.52 KB
Newer Older
Patrick von Platen's avatar
Patrick von Platen committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
# Copyright 2022 The HuggingFace Team. All rights reserved.
#
# 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.
import math

Patrick von Platen's avatar
Patrick von Platen committed
16
import numpy as np
Patrick von Platen's avatar
Patrick von Platen committed
17
18

from ..configuration_utils import ConfigMixin
Patrick von Platen's avatar
Patrick von Platen committed
19
from .scheduling_utils import SchedulerMixin, betas_for_alpha_bar, linear_beta_schedule
Patrick von Platen's avatar
Patrick von Platen committed
20
21


Patrick von Platen's avatar
Patrick von Platen committed
22
class DDIMScheduler(SchedulerMixin, ConfigMixin):
Patrick von Platen's avatar
Patrick von Platen committed
23
24
25
26
27
28
    def __init__(
        self,
        timesteps=1000,
        beta_start=0.0001,
        beta_end=0.02,
        beta_schedule="linear",
patil-suraj's avatar
patil-suraj committed
29
30
        trained_betas=None,
        timestep_values=None,
Patrick von Platen's avatar
Patrick von Platen committed
31
        clip_sample=True,
Patrick von Platen's avatar
Patrick von Platen committed
32
        tensor_format="np",
Patrick von Platen's avatar
Patrick von Platen committed
33
34
    ):
        super().__init__()
35
        self.register_to_config(
Patrick von Platen's avatar
Patrick von Platen committed
36
37
38
39
            timesteps=timesteps,
            beta_start=beta_start,
            beta_end=beta_end,
            beta_schedule=beta_schedule,
40
41
42
            trained_betas=trained_betas,
            timestep_values=timestep_values,
            clip_sample=clip_sample,
Patrick von Platen's avatar
Patrick von Platen committed
43
44
        )

45
        if beta_schedule == "linear":
Patrick von Platen's avatar
Patrick von Platen committed
46
            self.betas = linear_beta_schedule(timesteps, beta_start=beta_start, beta_end=beta_end)
Patrick von Platen's avatar
Patrick von Platen committed
47
48
        elif beta_schedule == "squaredcos_cap_v2":
            # GLIDE cosine schedule
Patrick von Platen's avatar
Patrick von Platen committed
49
            self.betas = betas_for_alpha_bar(
Patrick von Platen's avatar
Patrick von Platen committed
50
51
52
53
54
55
                timesteps,
                lambda t: math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2,
            )
        else:
            raise NotImplementedError(f"{beta_schedule} does is not implemented for {self.__class__}")

Patrick von Platen's avatar
Patrick von Platen committed
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
        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)

    #        alphas_cumprod_prev = torch.nn.functional.pad(alphas_cumprod[:-1], (1, 0), value=1.0)
    # TODO(PVP) - check how much of these is actually necessary!
    # LDM only uses "fixed_small"; glide seems to use a weird mix of the two, ...
    # https://github.com/openai/glide-text2im/blob/69b530740eb6cef69442d6180579ef5ba9ef063e/glide_text2im/gaussian_diffusion.py#L246
    #        variance = betas * (1.0 - alphas_cumprod_prev) / (1.0 - alphas_cumprod)
    #        if variance_type == "fixed_small":
    #            log_variance = torch.log(variance.clamp(min=1e-20))
    #        elif variance_type == "fixed_large":
    #            log_variance = torch.log(torch.cat([variance[1:2], betas[1:]], dim=0))
    #
    #
    #        self.register_buffer("log_variance", log_variance.to(torch.float32))
Patrick von Platen's avatar
Patrick von Platen committed
74

anton-l's avatar
anton-l committed
75
76
77
78
79
80
81
82
83
    # def rescale_betas(self, num_timesteps):
    #     # GLIDE scaling
    #     if self.beta_schedule == "linear":
    #         scale = self.timesteps / num_timesteps
    #         self.betas = linear_beta_schedule(
    #             num_timesteps, beta_start=self.beta_start * scale, beta_end=self.beta_end * scale
    #         )
    #         self.alphas = 1.0 - self.betas
    #         self.alphas_cumprod = np.cumprod(self.alphas, axis=0)
84
85
    def get_timestep_values(self):
        return self.config.timestep_values
anton-l's avatar
anton-l committed
86

Patrick von Platen's avatar
Patrick von Platen committed
87
88
89
90
91
92
93
94
    def get_alpha(self, time_step):
        return self.alphas[time_step]

    def get_beta(self, time_step):
        return self.betas[time_step]

    def get_alpha_prod(self, time_step):
        if time_step < 0:
Patrick von Platen's avatar
Patrick von Platen committed
95
            return self.one
Patrick von Platen's avatar
Patrick von Platen committed
96
97
        return self.alphas_cumprod[time_step]

Patrick von Platen's avatar
Patrick von Platen committed
98
99
100
    def get_orig_t(self, t, num_inference_steps):
        if t < 0:
            return -1
101
        return self.config.timesteps // num_inference_steps * t
Patrick von Platen's avatar
Patrick von Platen committed
102

Patrick von Platen's avatar
Patrick von Platen committed
103
    def get_variance(self, t, num_inference_steps):
Patrick von Platen's avatar
Patrick von Platen committed
104
105
        orig_t = self.get_orig_t(t, num_inference_steps)
        orig_prev_t = self.get_orig_t(t - 1, num_inference_steps)
Patrick von Platen's avatar
Patrick von Platen committed
106
107
108
109
110
111
112
113
114
115

        alpha_prod_t = self.get_alpha_prod(orig_t)
        alpha_prod_t_prev = self.get_alpha_prod(orig_prev_t)
        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

116
    def step(self, residual, sample, t, num_inference_steps, eta, use_clipped_residual=False):
Patrick von Platen's avatar
Patrick von Platen committed
117
118
119
120
121
        # 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)
122
        # - pred_original_sample -> f_theta(x_t, t) or x_0
Patrick von Platen's avatar
Patrick von Platen committed
123
124
        # - std_dev_t -> sigma_t
        # - eta -> η
125
126
        # - pred_sample_direction -> "direction pointingc to x_t"
        # - pred_prev_sample -> "x_t-1"
Patrick von Platen's avatar
Patrick von Platen committed
127
128

        # 1. get actual t and t-1
Patrick von Platen's avatar
Patrick von Platen committed
129
130
        orig_t = self.get_orig_t(t, num_inference_steps)
        orig_prev_t = self.get_orig_t(t - 1, num_inference_steps)
Patrick von Platen's avatar
Patrick von Platen committed
131
132
133
134
135
136

        # 2. compute alphas, betas
        alpha_prod_t = self.get_alpha_prod(orig_t)
        alpha_prod_t_prev = self.get_alpha_prod(orig_prev_t)
        beta_prod_t = 1 - alpha_prod_t

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

        # 4. Clip "predicted x_0"
142
        if self.config.clip_sample:
143
            pred_original_sample = self.clip(pred_original_sample, -1, 1)
Patrick von Platen's avatar
Patrick von Platen committed
144
145
146
147

        # 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
148
        std_dev_t = eta * variance ** (0.5)
Patrick von Platen's avatar
Patrick von Platen committed
149

anton-l's avatar
anton-l committed
150
151
        if use_clipped_residual:
            # the residual is always re-derived from the clipped x_0 in GLIDE
152
            residual = (sample - alpha_prod_t ** (0.5) * pred_original_sample) / beta_prod_t ** (0.5)
anton-l's avatar
anton-l committed
153

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

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

160
        return pred_prev_sample
Patrick von Platen's avatar
Patrick von Platen committed
161
162

    def __len__(self):
163
        return self.config.timesteps