scheduling_pndm.py 8.86 KB
Newer Older
1
# Copyright 2022 Zhejiang 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

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

17
import math
18
from typing import Union
Patrick von Platen's avatar
Patrick von Platen committed
19

20
import numpy as np
21
import torch
22

Patrick von Platen's avatar
Patrick von Platen committed
23
from ..configuration_utils import ConfigMixin
24
25
26
27
28
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
29
30
    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].
31

Patrick von Platen's avatar
Patrick von Platen committed
32
33
34
    :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.
35
36
37
    :param max_beta: the maximum beta to use; use values lower than 1 to
                     prevent singularities.
    """
38

39
40
41
42
43
44
45
46
47
    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
48
49
50
51
52


class PNDMScheduler(SchedulerMixin, ConfigMixin):
    def __init__(
        self,
Nathan Lambert's avatar
Nathan Lambert committed
53
        num_train_timesteps=1000,
Patrick von Platen's avatar
Patrick von Platen committed
54
55
56
57
58
59
        beta_start=0.0001,
        beta_end=0.02,
        beta_schedule="linear",
        tensor_format="np",
    ):
        super().__init__()
60
        self.register_to_config(
Nathan Lambert's avatar
Nathan Lambert committed
61
            num_train_timesteps=num_train_timesteps,
Patrick von Platen's avatar
Patrick von Platen committed
62
63
64
65
66
67
            beta_start=beta_start,
            beta_end=beta_end,
            beta_schedule=beta_schedule,
        )

        if beta_schedule == "linear":
Nathan Lambert's avatar
Nathan Lambert committed
68
            self.betas = np.linspace(beta_start, beta_end, num_train_timesteps, dtype=np.float32)
Patrick von Platen's avatar
Patrick von Platen committed
69
        elif beta_schedule == "squaredcos_cap_v2":
Patrick von Platen's avatar
Patrick von Platen committed
70
            # Glide cosine schedule
Nathan Lambert's avatar
Nathan Lambert committed
71
            self.betas = betas_for_alpha_bar(num_train_timesteps)
Patrick von Platen's avatar
Patrick von Platen committed
72
73
74
75
76
77
78
79
80
81
        else:
            raise NotImplementedError(f"{beta_schedule} does is not implemented for {self.__class__}")

        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
82
83
        # For now we only support F-PNDM, i.e. the runge-kutta method
        # For more information on the algorithm please take a look at the paper: https://arxiv.org/pdf/2202.09778.pdf
Patrick von Platen's avatar
Patrick von Platen committed
84
        # mainly at formula (9), (12), (13) and the Algorithm 2.
Patrick von Platen's avatar
Patrick von Platen committed
85
86
87
        self.pndm_order = 4

        # running values
Patrick von Platen's avatar
Patrick von Platen committed
88
        self.cur_model_output = 0
89
        self.cur_sample = None
Patrick von Platen's avatar
Patrick von Platen committed
90
        self.ets = []
Patrick von Platen's avatar
Patrick von Platen committed
91
        self.prk_time_steps = {}
Patrick von Platen's avatar
Patrick von Platen committed
92
        self.time_steps = {}
Patrick von Platen's avatar
Patrick von Platen committed
93
        self.set_prk_mode()
Patrick von Platen's avatar
Patrick von Platen committed
94

Patrick von Platen's avatar
Patrick von Platen committed
95
96
97
    def get_prk_time_steps(self, num_inference_steps):
        if num_inference_steps in self.prk_time_steps:
            return self.prk_time_steps[num_inference_steps]
Patrick von Platen's avatar
Patrick von Platen committed
98

Nathan Lambert's avatar
Nathan Lambert committed
99
100
101
        inference_step_times = list(
            range(0, self.config.num_train_timesteps, self.config.num_train_timesteps // num_inference_steps)
        )
Patrick von Platen's avatar
Patrick von Platen committed
102

Patrick von Platen's avatar
Patrick von Platen committed
103
        prk_time_steps = np.array(inference_step_times[-self.pndm_order :]).repeat(2) + np.tile(
Nathan Lambert's avatar
Nathan Lambert committed
104
            np.array([0, self.config.num_train_timesteps // num_inference_steps // 2]), self.pndm_order
105
        )
Patrick von Platen's avatar
Patrick von Platen committed
106
        self.prk_time_steps[num_inference_steps] = list(reversed(prk_time_steps[:-1].repeat(2)[1:-1]))
Patrick von Platen's avatar
Patrick von Platen committed
107

Patrick von Platen's avatar
Patrick von Platen committed
108
        return self.prk_time_steps[num_inference_steps]
Patrick von Platen's avatar
Patrick von Platen committed
109

Patrick von Platen's avatar
Patrick von Platen committed
110
111
112
    def get_time_steps(self, num_inference_steps):
        if num_inference_steps in self.time_steps:
            return self.time_steps[num_inference_steps]
Patrick von Platen's avatar
Patrick von Platen committed
113

Nathan Lambert's avatar
Nathan Lambert committed
114
115
116
        inference_step_times = list(
            range(0, self.config.num_train_timesteps, self.config.num_train_timesteps // num_inference_steps)
        )
Patrick von Platen's avatar
Patrick von Platen committed
117
        self.time_steps[num_inference_steps] = list(reversed(inference_step_times[:-3]))
Patrick von Platen's avatar
Patrick von Platen committed
118

Patrick von Platen's avatar
Patrick von Platen committed
119
        return self.time_steps[num_inference_steps]
Patrick von Platen's avatar
Patrick von Platen committed
120

Patrick von Platen's avatar
Patrick von Platen committed
121
122
123
124
125
126
127
128
129
130
131
132
133
134
    def set_prk_mode(self):
        self.mode = "prk"

    def set_plms_mode(self):
        self.mode = "plms"

    def step(self, *args, **kwargs):
        if self.mode == "prk":
            return self.step_prk(*args, **kwargs)
        if self.mode == "plms":
            return self.step_plms(*args, **kwargs)

        raise ValueError(f"mode {self.mode} does not exist.")

135
136
    def step_prk(
        self,
Patrick von Platen's avatar
Patrick von Platen committed
137
        model_output: Union[torch.FloatTensor, np.ndarray],
138
139
140
141
        timestep: int,
        sample: Union[torch.FloatTensor, np.ndarray],
        num_inference_steps,
    ):
Nathan Lambert's avatar
Nathan Lambert committed
142
143
144
145
        """
        Step function propagating the sample with the Runge-Kutta method. RK takes 4 forward passes to approximate the
        solution to the differential equation.
        """
146
        t = timestep
Patrick von Platen's avatar
Patrick von Platen committed
147
        prk_time_steps = self.get_prk_time_steps(num_inference_steps)
Patrick von Platen's avatar
Patrick von Platen committed
148

Patrick von Platen's avatar
Patrick von Platen committed
149
150
        t_orig = prk_time_steps[t // 4 * 4]
        t_orig_prev = prk_time_steps[min(t + 1, len(prk_time_steps) - 1)]
Patrick von Platen's avatar
Patrick von Platen committed
151

Patrick von Platen's avatar
Patrick von Platen committed
152
        if t % 4 == 0:
Patrick von Platen's avatar
Patrick von Platen committed
153
154
            self.cur_model_output += 1 / 6 * model_output
            self.ets.append(model_output)
155
            self.cur_sample = sample
Patrick von Platen's avatar
Patrick von Platen committed
156
        elif (t - 1) % 4 == 0:
Patrick von Platen's avatar
Patrick von Platen committed
157
            self.cur_model_output += 1 / 3 * model_output
Patrick von Platen's avatar
Patrick von Platen committed
158
        elif (t - 2) % 4 == 0:
Patrick von Platen's avatar
Patrick von Platen committed
159
            self.cur_model_output += 1 / 3 * model_output
Patrick von Platen's avatar
Patrick von Platen committed
160
        elif (t - 3) % 4 == 0:
Patrick von Platen's avatar
Patrick von Platen committed
161
162
            model_output = self.cur_model_output + 1 / 6 * model_output
            self.cur_model_output = 0
Patrick von Platen's avatar
Patrick von Platen committed
163

Patrick von Platen's avatar
Patrick von Platen committed
164
165
166
        # cur_sample should not be `None`
        cur_sample = self.cur_sample if self.cur_sample is not None else sample

Patrick von Platen's avatar
Patrick von Platen committed
167
        return {"prev_sample": self.get_prev_sample(cur_sample, t_orig, t_orig_prev, model_output)}
Patrick von Platen's avatar
Patrick von Platen committed
168

169
170
    def step_plms(
        self,
Patrick von Platen's avatar
Patrick von Platen committed
171
        model_output: Union[torch.FloatTensor, np.ndarray],
172
173
174
175
        timestep: int,
        sample: Union[torch.FloatTensor, np.ndarray],
        num_inference_steps,
    ):
Nathan Lambert's avatar
Nathan Lambert committed
176
177
178
179
        """
        Step function propagating the sample with the linear multi-step method. This has one forward pass with multiple
        times to approximate the solution.
        """
180
        t = timestep
Patrick von Platen's avatar
Patrick von Platen committed
181
182
183
184
185
186
187
188
        if len(self.ets) < 3:
            raise ValueError(
                f"{self.__class__} can only be run AFTER scheduler has been run "
                "in 'prk' mode for at least 12 iterations "
                "See: https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/pipeline_pndm.py "
                "for more information."
            )

Patrick von Platen's avatar
Patrick von Platen committed
189
190
        timesteps = self.get_time_steps(num_inference_steps)

Patrick von Platen's avatar
Patrick von Platen committed
191
192
        t_orig = timesteps[t]
        t_orig_prev = timesteps[min(t + 1, len(timesteps) - 1)]
Patrick von Platen's avatar
Patrick von Platen committed
193
        self.ets.append(model_output)
Patrick von Platen's avatar
Patrick von Platen committed
194

Patrick von Platen's avatar
Patrick von Platen committed
195
        model_output = (1 / 24) * (55 * self.ets[-1] - 59 * self.ets[-2] + 37 * self.ets[-3] - 9 * self.ets[-4])
Patrick von Platen's avatar
Patrick von Platen committed
196

Patrick von Platen's avatar
Patrick von Platen committed
197
        return {"prev_sample": self.get_prev_sample(sample, t_orig, t_orig_prev, model_output)}
Patrick von Platen's avatar
Patrick von Platen committed
198

Patrick von Platen's avatar
Patrick von Platen committed
199
    def get_prev_sample(self, sample, t_orig, t_orig_prev, model_output):
Patrick von Platen's avatar
Patrick von Platen committed
200
201
202
203
204
205
206
207
208
209
        # See formula (9) of PNDM paper https://arxiv.org/pdf/2202.09778.pdf
        # this function computes x_(t−δ) using the formula of (9)
        # Note that x_t needs to be added to both sides of the equation

        # Notation (<variable name> -> <name in paper>
        # alpha_prod_t -> α_t
        # alpha_prod_t_prev -> α_(t−δ)
        # beta_prod_t -> (1 - α_t)
        # beta_prod_t_prev -> (1 - α_(t−δ))
        # sample -> x_t
Patrick von Platen's avatar
Patrick von Platen committed
210
        # model_output -> e_θ(x_t, t)
Patrick von Platen's avatar
Patrick von Platen committed
211
        # prev_sample -> x_(t−δ)
212
213
        alpha_prod_t = self.alphas_cumprod[t_orig + 1]
        alpha_prod_t_prev = self.alphas_cumprod[t_orig_prev + 1]
Patrick von Platen's avatar
Patrick von Platen committed
214
215
216
217
218
219
220
221
222
223
        beta_prod_t = 1 - alpha_prod_t
        beta_prod_t_prev = 1 - alpha_prod_t_prev

        # corresponds to (α_(t−δ) - α_t) divided by
        # denominator of x_t in formula (9) and plus 1
        # Note: (α_(t−δ) - α_t) / (sqrt(α_t) * (sqrt(α_(t−δ)) + sqr(α_t))) =
        # sqrt(α_(t−δ)) / sqrt(α_t))
        sample_coeff = (alpha_prod_t_prev / alpha_prod_t) ** (0.5)

        # corresponds to denominator of e_θ(x_t, t) in formula (9)
Patrick von Platen's avatar
Patrick von Platen committed
224
        model_output_denom_coeff = alpha_prod_t * beta_prod_t_prev ** (0.5) + (
Patrick von Platen's avatar
Patrick von Platen committed
225
226
227
228
            alpha_prod_t * beta_prod_t * alpha_prod_t_prev
        ) ** (0.5)

        # full formula (9)
Patrick von Platen's avatar
Patrick von Platen committed
229
230
231
        prev_sample = (
            sample_coeff * sample - (alpha_prod_t_prev - alpha_prod_t) * model_output / model_output_denom_coeff
        )
Patrick von Platen's avatar
Patrick von Platen committed
232
233

        return prev_sample
Patrick von Platen's avatar
Patrick von Platen committed
234
235

    def __len__(self):
Nathan Lambert's avatar
Nathan Lambert committed
236
        return self.config.num_train_timesteps