scheduling_pndm.py 8.31 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
53
54
55
56
57
58
59


class PNDMScheduler(SchedulerMixin, ConfigMixin):
    def __init__(
        self,
        timesteps=1000,
        beta_start=0.0001,
        beta_end=0.02,
        beta_schedule="linear",
        tensor_format="np",
    ):
        super().__init__()
60
        self.register_to_config(
Patrick von Platen's avatar
Patrick von Platen committed
61
62
63
64
65
66
67
            timesteps=timesteps,
            beta_start=beta_start,
            beta_end=beta_end,
            beta_schedule=beta_schedule,
        )

        if beta_schedule == "linear":
68
            self.betas = np.linspace(beta_start, beta_end, 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
71
            self.betas = betas_for_alpha_bar(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

99
        inference_step_times = list(range(0, self.config.timesteps, self.config.timesteps // num_inference_steps))
Patrick von Platen's avatar
Patrick von Platen committed
100

Patrick von Platen's avatar
Patrick von Platen committed
101
        prk_time_steps = np.array(inference_step_times[-self.pndm_order :]).repeat(2) + np.tile(
102
            np.array([0, self.config.timesteps // num_inference_steps // 2]), self.pndm_order
103
        )
Patrick von Platen's avatar
Patrick von Platen committed
104
        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
105

Patrick von Platen's avatar
Patrick von Platen committed
106
        return self.prk_time_steps[num_inference_steps]
Patrick von Platen's avatar
Patrick von Platen committed
107

Patrick von Platen's avatar
Patrick von Platen committed
108
109
110
    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
111

112
        inference_step_times = list(range(0, self.config.timesteps, self.config.timesteps // num_inference_steps))
Patrick von Platen's avatar
Patrick von Platen committed
113
        self.time_steps[num_inference_steps] = list(reversed(inference_step_times[:-3]))
Patrick von Platen's avatar
Patrick von Platen committed
114

Patrick von Platen's avatar
Patrick von Platen committed
115
        return self.time_steps[num_inference_steps]
Patrick von Platen's avatar
Patrick von Platen committed
116

Patrick von Platen's avatar
Patrick von Platen committed
117
118
119
120
121
122
123
124
125
126
127
128
129
130
    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.")

131
132
    def step_prk(
        self,
Patrick von Platen's avatar
Patrick von Platen committed
133
        model_output: Union[torch.FloatTensor, np.ndarray],
134
135
136
137
138
        timestep: int,
        sample: Union[torch.FloatTensor, np.ndarray],
        num_inference_steps,
    ):
        t = timestep
Patrick von Platen's avatar
Patrick von Platen committed
139
        prk_time_steps = self.get_prk_time_steps(num_inference_steps)
Patrick von Platen's avatar
Patrick von Platen committed
140

Patrick von Platen's avatar
Patrick von Platen committed
141
142
        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
143

Patrick von Platen's avatar
Patrick von Platen committed
144
        if t % 4 == 0:
Patrick von Platen's avatar
Patrick von Platen committed
145
146
            self.cur_model_output += 1 / 6 * model_output
            self.ets.append(model_output)
147
            self.cur_sample = sample
Patrick von Platen's avatar
Patrick von Platen committed
148
        elif (t - 1) % 4 == 0:
Patrick von Platen's avatar
Patrick von Platen committed
149
            self.cur_model_output += 1 / 3 * model_output
Patrick von Platen's avatar
Patrick von Platen committed
150
        elif (t - 2) % 4 == 0:
Patrick von Platen's avatar
Patrick von Platen committed
151
            self.cur_model_output += 1 / 3 * model_output
Patrick von Platen's avatar
Patrick von Platen committed
152
        elif (t - 3) % 4 == 0:
Patrick von Platen's avatar
Patrick von Platen committed
153
154
            model_output = self.cur_model_output + 1 / 6 * model_output
            self.cur_model_output = 0
Patrick von Platen's avatar
Patrick von Platen committed
155

Patrick von Platen's avatar
Patrick von Platen committed
156
157
158
        # 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
159
        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
160

161
162
    def step_plms(
        self,
Patrick von Platen's avatar
Patrick von Platen committed
163
        model_output: Union[torch.FloatTensor, np.ndarray],
164
165
166
167
168
        timestep: int,
        sample: Union[torch.FloatTensor, np.ndarray],
        num_inference_steps,
    ):
        t = timestep
Patrick von Platen's avatar
Patrick von Platen committed
169
170
171
172
173
174
175
176
        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
177
178
        timesteps = self.get_time_steps(num_inference_steps)

Patrick von Platen's avatar
Patrick von Platen committed
179
180
        t_orig = timesteps[t]
        t_orig_prev = timesteps[min(t + 1, len(timesteps) - 1)]
Patrick von Platen's avatar
Patrick von Platen committed
181
        self.ets.append(model_output)
Patrick von Platen's avatar
Patrick von Platen committed
182

Patrick von Platen's avatar
Patrick von Platen committed
183
        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
184

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

Patrick von Platen's avatar
Patrick von Platen committed
187
    def get_prev_sample(self, sample, t_orig, t_orig_prev, model_output):
Patrick von Platen's avatar
Patrick von Platen committed
188
189
190
191
192
193
194
195
196
197
        # 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
198
        # model_output -> e_θ(x_t, t)
Patrick von Platen's avatar
Patrick von Platen committed
199
        # prev_sample -> x_(t−δ)
200
201
        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
202
203
204
205
206
207
208
209
210
211
        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
212
        model_output_denom_coeff = alpha_prod_t * beta_prod_t_prev ** (0.5) + (
Patrick von Platen's avatar
Patrick von Platen committed
213
214
215
216
            alpha_prod_t * beta_prod_t * alpha_prod_t_prev
        ) ** (0.5)

        # full formula (9)
Patrick von Platen's avatar
Patrick von Platen committed
217
        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
218
219

        return prev_sample
Patrick von Platen's avatar
Patrick von Platen committed
220
221

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