scheduling_ddpm.py 23.6 KB
Newer Older
Patrick von Platen's avatar
Patrick von Platen committed
1
# Copyright 2023 UC Berkeley 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
18
from dataclasses import dataclass
19
from typing import List, Optional, Tuple, Union
Patrick von Platen's avatar
Patrick von Platen committed
20

Patrick von Platen's avatar
Patrick von Platen committed
21
import numpy as np
22
import torch
Patrick von Platen's avatar
improve  
Patrick von Platen committed
23

24
from ..configuration_utils import ConfigMixin, register_to_config
Dhruv Nair's avatar
Dhruv Nair committed
25
26
from ..utils import BaseOutput
from ..utils.torch_utils import randn_tensor
Kashif Rasul's avatar
Kashif Rasul committed
27
from .scheduling_utils import KarrasDiffusionSchedulers, SchedulerMixin
28
29
30
31
32


@dataclass
class DDPMSchedulerOutput(BaseOutput):
    """
33
    Output class for the scheduler's `step` function output.
34
35
36

    Args:
        prev_sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` for images):
37
            Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next model input in the
38
39
            denoising loop.
        pred_original_sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` for images):
40
            The predicted denoised sample `(x_{0})` based on the model output from the current timestep.
41
42
43
44
45
            `pred_original_sample` can be used to preview progress or for guidance.
    """

    prev_sample: torch.FloatTensor
    pred_original_sample: Optional[torch.FloatTensor] = None
46
47


YiYi Xu's avatar
YiYi Xu committed
48
49
50
51
52
def betas_for_alpha_bar(
    num_diffusion_timesteps,
    max_beta=0.999,
    alpha_transform_type="cosine",
):
53
    """
Patrick von Platen's avatar
Patrick von Platen committed
54
55
    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].
56

57
58
59
60
61
62
63
    Contains a function alpha_bar that takes an argument t and transforms it to the cumulative product of (1-beta) up
    to that part of the diffusion process.


    Args:
        num_diffusion_timesteps (`int`): the number of betas to produce.
        max_beta (`float`): the maximum beta to use; use values lower than 1 to
64
                     prevent singularities.
YiYi Xu's avatar
YiYi Xu committed
65
66
        alpha_transform_type (`str`, *optional*, default to `cosine`): the type of noise schedule for alpha_bar.
                     Choose from `cosine` or `exp`
67
68
69

    Returns:
        betas (`np.ndarray`): the betas used by the scheduler to step the model outputs
70
    """
YiYi Xu's avatar
YiYi Xu committed
71
    if alpha_transform_type == "cosine":
72

YiYi Xu's avatar
YiYi Xu committed
73
74
75
76
77
78
79
80
81
82
        def alpha_bar_fn(t):
            return math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2

    elif alpha_transform_type == "exp":

        def alpha_bar_fn(t):
            return math.exp(t * -12.0)

    else:
        raise ValueError(f"Unsupported alpha_tranform_type: {alpha_transform_type}")
83
84
85
86
87

    betas = []
    for i in range(num_diffusion_timesteps):
        t1 = i / num_diffusion_timesteps
        t2 = (i + 1) / num_diffusion_timesteps
YiYi Xu's avatar
YiYi Xu committed
88
        betas.append(min(1 - alpha_bar_fn(t2) / alpha_bar_fn(t1), max_beta))
89
    return torch.tensor(betas, dtype=torch.float32)
Patrick von Platen's avatar
improve  
Patrick von Platen committed
90
91


Patrick von Platen's avatar
Patrick von Platen committed
92
class DDPMScheduler(SchedulerMixin, ConfigMixin):
93
    """
94
    `DDPMScheduler` explores the connections between denoising score matching and Langevin dynamics sampling.
95

96
97
    This model inherits from [`SchedulerMixin`] and [`ConfigMixin`]. Check the superclass documentation for the generic
    methods the library implements for all schedulers such as loading and saving.
98
99

    Args:
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
        num_train_timesteps (`int`, defaults to 1000):
            The number of diffusion steps to train the model.
        beta_start (`float`, defaults to 0.0001):
            The starting `beta` value of inference.
        beta_end (`float`, defaults to 0.02):
            The final `beta` value.
        beta_schedule (`str`, defaults to `"linear"`):
            The beta schedule, a mapping from a beta range to a sequence of betas for stepping the model. Choose from
            `linear`, `scaled_linear`, or `squaredcos_cap_v2`.
        variance_type (`str`, defaults to `"fixed_small"`):
            Clip the variance when adding noise to the denoised sample. Choose from `fixed_small`, `fixed_small_log`,
            `fixed_large`, `fixed_large_log`, `learned` or `learned_range`.
        clip_sample (`bool`, defaults to `True`):
            Clip the predicted sample for numerical stability.
        clip_sample_range (`float`, defaults to 1.0):
            The maximum magnitude for sample clipping. Valid only when `clip_sample=True`.
        prediction_type (`str`, defaults to `epsilon`, *optional*):
            Prediction type of the scheduler function; can be `epsilon` (predicts the noise of the diffusion process),
            `sample` (directly predicts the noisy sample`) or `v_prediction` (see section 2.4 of [Imagen
            Video](https://imagen.research.google/video/paper.pdf) paper).
        thresholding (`bool`, defaults to `False`):
            Whether to use the "dynamic thresholding" method. This is unsuitable for latent-space diffusion models such
            as Stable Diffusion.
        dynamic_thresholding_ratio (`float`, defaults to 0.995):
            The ratio for the dynamic thresholding method. Valid only when `thresholding=True`.
        sample_max_value (`float`, defaults to 1.0):
            The threshold value for dynamic thresholding. Valid only when `thresholding=True`.
        timestep_spacing (`str`, defaults to `"leading"`):
            The way the timesteps should be scaled. Refer to Table 2 of the [Common Diffusion Noise Schedules and
            Sample Steps are Flawed](https://huggingface.co/papers/2305.08891) for more information.
        steps_offset (`int`, defaults to 0):
            An offset added to the inference steps. You can use a combination of `offset=1` and
            `set_alpha_to_one=False` to make the last step use step 0 for the previous alpha product like in Stable
            Diffusion.
134
135
    """

Kashif Rasul's avatar
Kashif Rasul committed
136
    _compatibles = [e.name for e in KarrasDiffusionSchedulers]
137
    order = 1
138

139
    @register_to_config
Patrick von Platen's avatar
improve  
Patrick von Platen committed
140
141
    def __init__(
        self,
Partho's avatar
Partho committed
142
143
144
145
        num_train_timesteps: int = 1000,
        beta_start: float = 0.0001,
        beta_end: float = 0.02,
        beta_schedule: str = "linear",
146
        trained_betas: Optional[Union[np.ndarray, List[float]]] = None,
Partho's avatar
Partho committed
147
148
        variance_type: str = "fixed_small",
        clip_sample: bool = True,
149
        prediction_type: str = "epsilon",
150
151
152
153
        thresholding: bool = False,
        dynamic_thresholding_ratio: float = 0.995,
        clip_sample_range: float = 1.0,
        sample_max_value: float = 1.0,
154
155
        timestep_spacing: str = "leading",
        steps_offset: int = 0,
Patrick von Platen's avatar
improve  
Patrick von Platen committed
156
    ):
157
        if trained_betas is not None:
158
            self.betas = torch.tensor(trained_betas, dtype=torch.float32)
159
        elif beta_schedule == "linear":
160
            self.betas = torch.linspace(beta_start, beta_end, num_train_timesteps, dtype=torch.float32)
161
162
        elif beta_schedule == "scaled_linear":
            # this schedule is very specific to the latent diffusion model.
163
            self.betas = torch.linspace(beta_start**0.5, beta_end**0.5, num_train_timesteps, dtype=torch.float32) ** 2
anton-l's avatar
anton-l committed
164
        elif beta_schedule == "squaredcos_cap_v2":
Patrick von Platen's avatar
Patrick von Platen committed
165
            # Glide cosine schedule
Nathan Lambert's avatar
Nathan Lambert committed
166
            self.betas = betas_for_alpha_bar(num_train_timesteps)
Nathan Lambert's avatar
Nathan Lambert committed
167
168
169
170
        elif beta_schedule == "sigmoid":
            # GeoDiff sigmoid schedule
            betas = torch.linspace(-6, 6, num_train_timesteps)
            self.betas = torch.sigmoid(betas) * (beta_end - beta_start) + beta_start
Patrick von Platen's avatar
improve  
Patrick von Platen committed
171
172
173
        else:
            raise NotImplementedError(f"{beta_schedule} does is not implemented for {self.__class__}")

Patrick von Platen's avatar
Patrick von Platen committed
174
        self.alphas = 1.0 - self.betas
175
176
        self.alphas_cumprod = torch.cumprod(self.alphas, dim=0)
        self.one = torch.tensor(1.0)
Patrick von Platen's avatar
Patrick von Platen committed
177

178
179
180
        # standard deviation of the initial noise distribution
        self.init_noise_sigma = 1.0

181
        # setable values
Will Berman's avatar
Will Berman committed
182
        self.custom_timesteps = False
183
        self.num_inference_steps = None
184
        self.timesteps = torch.from_numpy(np.arange(0, num_train_timesteps)[::-1].copy())
Patrick von Platen's avatar
Patrick von Platen committed
185

186
187
        self.variance_type = variance_type

188
189
190
191
192
193
    def scale_model_input(self, sample: torch.FloatTensor, timestep: Optional[int] = None) -> torch.FloatTensor:
        """
        Ensures interchangeability with schedulers that need to scale the denoising model input depending on the
        current timestep.

        Args:
194
195
196
197
            sample (`torch.FloatTensor`):
                The input sample.
            timestep (`int`, *optional*):
                The current timestep in the diffusion chain.
198
199

        Returns:
200
201
            `torch.FloatTensor`:
                A scaled input sample.
202
203
204
        """
        return sample

Will Berman's avatar
Will Berman committed
205
206
207
208
209
210
    def set_timesteps(
        self,
        num_inference_steps: Optional[int] = None,
        device: Union[str, torch.device] = None,
        timesteps: Optional[List[int]] = None,
    ):
211
        """
212
        Sets the discrete timesteps used for the diffusion chain (to be run before inference).
213
214

        Args:
215
216
            num_inference_steps (`int`):
                The number of diffusion steps used when generating samples with a pre-trained model. If used,
Will Berman's avatar
Will Berman committed
217
                `timesteps` must be `None`.
218
219
220
221
222
223
            device (`str` or `torch.device`, *optional*):
                The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
            timesteps (`List[int]`, *optional*):
                Custom timesteps used to support arbitrary spacing between timesteps. If `None`, then the default
                timestep spacing strategy of equal spacing between timesteps is used. If `timesteps` is passed,
                `num_inference_steps` must be `None`.
Will Berman's avatar
Will Berman committed
224

225
        """
Will Berman's avatar
Will Berman committed
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
        if num_inference_steps is not None and timesteps is not None:
            raise ValueError("Can only pass one of `num_inference_steps` or `custom_timesteps`.")

        if timesteps is not None:
            for i in range(1, len(timesteps)):
                if timesteps[i] >= timesteps[i - 1]:
                    raise ValueError("`custom_timesteps` must be in descending order.")

            if timesteps[0] >= self.config.num_train_timesteps:
                raise ValueError(
                    f"`timesteps` must start before `self.config.train_timesteps`:"
                    f" {self.config.num_train_timesteps}."
                )

            timesteps = np.array(timesteps, dtype=np.int64)
            self.custom_timesteps = True
        else:
            if num_inference_steps > self.config.num_train_timesteps:
                raise ValueError(
                    f"`num_inference_steps`: {num_inference_steps} cannot be larger than `self.config.train_timesteps`:"
                    f" {self.config.num_train_timesteps} as the unet model trained with this scheduler can only handle"
                    f" maximal {self.config.num_train_timesteps} timesteps."
                )
249

Will Berman's avatar
Will Berman committed
250
251
            self.num_inference_steps = num_inference_steps
            self.custom_timesteps = False
252

253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
            # "linspace", "leading", "trailing" corresponds to annotation of Table 2. of https://arxiv.org/abs/2305.08891
            if self.config.timestep_spacing == "linspace":
                timesteps = (
                    np.linspace(0, self.config.num_train_timesteps - 1, num_inference_steps)
                    .round()[::-1]
                    .copy()
                    .astype(np.int64)
                )
            elif self.config.timestep_spacing == "leading":
                step_ratio = self.config.num_train_timesteps // self.num_inference_steps
                # creates integer timesteps by multiplying by ratio
                # casting to int to avoid issues when num_inference_step is power of 3
                timesteps = (np.arange(0, num_inference_steps) * step_ratio).round()[::-1].copy().astype(np.int64)
                timesteps += self.config.steps_offset
            elif self.config.timestep_spacing == "trailing":
                step_ratio = self.config.num_train_timesteps / self.num_inference_steps
                # creates integer timesteps by multiplying by ratio
                # casting to int to avoid issues when num_inference_step is power of 3
                timesteps = np.round(np.arange(self.config.num_train_timesteps, 0, -step_ratio)).astype(np.int64)
                timesteps -= 1
            else:
                raise ValueError(
                    f"{self.config.timestep_spacing} is not supported. Please make sure to choose one of 'linspace', 'leading' or 'trailing'."
                )

278
        self.timesteps = torch.from_numpy(timesteps).to(device)
279

280
    def _get_variance(self, t, predicted_variance=None, variance_type=None):
Will Berman's avatar
Will Berman committed
281
282
        prev_t = self.previous_timestep(t)

283
        alpha_prod_t = self.alphas_cumprod[t]
284
285
        alpha_prod_t_prev = self.alphas_cumprod[prev_t] if prev_t >= 0 else self.one
        current_beta_t = 1 - alpha_prod_t / alpha_prod_t_prev
Patrick von Platen's avatar
Patrick von Platen committed
286

Kashif Rasul's avatar
Kashif Rasul committed
287
        # For t > 0, compute predicted variance βt (see formula (6) and (7) from https://arxiv.org/pdf/2006.11239.pdf)
288
        # and sample from it to get previous sample
Kashif Rasul's avatar
Kashif Rasul committed
289
        # x_{t-1} ~ N(pred_prev_sample, variance) == add variance to pred_sample
290
        variance = (1 - alpha_prod_t_prev) / (1 - alpha_prod_t) * current_beta_t
William Berman's avatar
William Berman committed
291
292

        # we always take the log of variance, so clamp it to ensure it's not 0
William Berman's avatar
William Berman committed
293
        variance = torch.clamp(variance, min=1e-20)
Patrick von Platen's avatar
Patrick von Platen committed
294

295
296
297
        if variance_type is None:
            variance_type = self.config.variance_type

298
        # hacks - were probably added for training stability
299
        if variance_type == "fixed_small":
William Berman's avatar
William Berman committed
300
            variance = variance
301
        # for rl-diffuser https://arxiv.org/abs/2205.09991
302
        elif variance_type == "fixed_small_log":
303
            variance = torch.log(variance)
304
            variance = torch.exp(0.5 * variance)
305
        elif variance_type == "fixed_large":
306
            variance = current_beta_t
307
        elif variance_type == "fixed_large_log":
Patrick von Platen's avatar
Patrick von Platen committed
308
            # Glide max_log
309
            variance = torch.log(current_beta_t)
310
311
312
        elif variance_type == "learned":
            return predicted_variance
        elif variance_type == "learned_range":
313
            min_log = torch.log(variance)
William Berman's avatar
William Berman committed
314
            max_log = torch.log(current_beta_t)
315
316
            frac = (predicted_variance + 1) / 2
            variance = frac * max_log + (1 - frac) * min_log
Patrick von Platen's avatar
Patrick von Platen committed
317
318
319

        return variance

320
    def _threshold_sample(self, sample: torch.FloatTensor) -> torch.FloatTensor:
321
322
323
324
325
326
327
328
329
330
        """
        "Dynamic thresholding: At each sampling step we set s to a certain percentile absolute pixel value in xt0 (the
        prediction of x_0 at timestep t), and if s > 1, then we threshold xt0 to the range [-s, s] and then divide by
        s. Dynamic thresholding pushes saturated pixels (those near -1 and 1) inwards, thereby actively preventing
        pixels from saturation at each step. We find that dynamic thresholding results in significantly better
        photorealism as well as better image-text alignment, especially when using very large guidance weights."

        https://arxiv.org/abs/2205.11487
        """
        dtype = sample.dtype
331
        batch_size, channels, *remaining_dims = sample.shape
332
333
334
335
336

        if dtype not in (torch.float32, torch.float64):
            sample = sample.float()  # upcast for quantile calculation, and clamp not implemented for cpu half

        # Flatten sample for doing quantile calculation along each image
337
        sample = sample.reshape(batch_size, channels * np.prod(remaining_dims))
338
339
340
341
342
343
344
345
346
347

        abs_sample = sample.abs()  # "a certain percentile absolute pixel value"

        s = torch.quantile(abs_sample, self.config.dynamic_thresholding_ratio, dim=1)
        s = torch.clamp(
            s, min=1, max=self.config.sample_max_value
        )  # When clamped to min=1, equivalent to standard clipping to [-1, 1]
        s = s.unsqueeze(1)  # (batch_size, 1) because clamp will broadcast along dim=0
        sample = torch.clamp(sample, -s, s) / s  # "we threshold xt0 to the range [-s, s] and then divide by s"

348
        sample = sample.reshape(batch_size, channels, *remaining_dims)
349
350
351
        sample = sample.to(dtype)

        return sample
352

353
354
    def step(
        self,
355
        model_output: torch.FloatTensor,
356
        timestep: int,
357
        sample: torch.FloatTensor,
Patrick von Platen's avatar
Patrick von Platen committed
358
        generator=None,
359
        return_dict: bool = True,
360
    ) -> Union[DDPMSchedulerOutput, Tuple]:
361
        """
362
        Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion
363
364
365
        process from the learned model outputs (most often the predicted noise).

        Args:
366
367
368
369
            model_output (`torch.FloatTensor`):
                The direct output from learned diffusion model.
            timestep (`float`):
                The current discrete timestep in the diffusion chain.
370
            sample (`torch.FloatTensor`):
371
372
373
374
375
                A current instance of a sample created by the diffusion process.
            generator (`torch.Generator`, *optional*):
                A random number generator.
            return_dict (`bool`, *optional*, defaults to `True`):
                Whether or not to return a [`~schedulers.scheduling_ddpm.DDPMSchedulerOutput`] or `tuple`.
376
377

        Returns:
378
379
380
            [`~schedulers.scheduling_ddpm.DDPMSchedulerOutput`] or `tuple`:
                If return_dict is `True`, [`~schedulers.scheduling_ddpm.DDPMSchedulerOutput`] is returned, otherwise a
                tuple is returned where the first element is the sample tensor.
381
382

        """
383
        t = timestep
Will Berman's avatar
Will Berman committed
384
385

        prev_t = self.previous_timestep(t)
386

387
388
389
390
391
        if model_output.shape[1] == sample.shape[1] * 2 and self.variance_type in ["learned", "learned_range"]:
            model_output, predicted_variance = torch.split(model_output, sample.shape[1], dim=1)
        else:
            predicted_variance = None

Patrick von Platen's avatar
Patrick von Platen committed
392
        # 1. compute alphas, betas
393
        alpha_prod_t = self.alphas_cumprod[t]
394
        alpha_prod_t_prev = self.alphas_cumprod[prev_t] if prev_t >= 0 else self.one
Patrick von Platen's avatar
Patrick von Platen committed
395
396
        beta_prod_t = 1 - alpha_prod_t
        beta_prod_t_prev = 1 - alpha_prod_t_prev
397
398
        current_alpha_t = alpha_prod_t / alpha_prod_t_prev
        current_beta_t = 1 - current_alpha_t
Patrick von Platen's avatar
Patrick von Platen committed
399

400
        # 2. compute predicted original sample from predicted noise also called
Patrick von Platen's avatar
Patrick von Platen committed
401
        # "predicted x_0" of formula (15) from https://arxiv.org/pdf/2006.11239.pdf
402
        if self.config.prediction_type == "epsilon":
Patrick von Platen's avatar
Patrick von Platen committed
403
            pred_original_sample = (sample - beta_prod_t ** (0.5) * model_output) / alpha_prod_t ** (0.5)
404
        elif self.config.prediction_type == "sample":
Patrick von Platen's avatar
Patrick von Platen committed
405
            pred_original_sample = model_output
406
407
        elif self.config.prediction_type == "v_prediction":
            pred_original_sample = (alpha_prod_t**0.5) * sample - (beta_prod_t**0.5) * model_output
408
409
        else:
            raise ValueError(
410
411
                f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample` or"
                " `v_prediction`  for the DDPMScheduler."
412
            )
Patrick von Platen's avatar
Patrick von Platen committed
413

414
        # 3. Clip or threshold "predicted x_0"
415
416
417
        if self.config.thresholding:
            pred_original_sample = self._threshold_sample(pred_original_sample)
        elif self.config.clip_sample:
418
419
            pred_original_sample = pred_original_sample.clamp(
                -self.config.clip_sample_range, self.config.clip_sample_range
Will Berman's avatar
Will Berman committed
420
            )
Patrick von Platen's avatar
Patrick von Platen committed
421

422
        # 4. Compute coefficients for pred_original_sample x_0 and current sample x_t
Patrick von Platen's avatar
Patrick von Platen committed
423
        # See formula (7) from https://arxiv.org/pdf/2006.11239.pdf
424
425
        pred_original_sample_coeff = (alpha_prod_t_prev ** (0.5) * current_beta_t) / beta_prod_t
        current_sample_coeff = current_alpha_t ** (0.5) * beta_prod_t_prev / beta_prod_t
Patrick von Platen's avatar
Patrick von Platen committed
426

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

Patrick von Platen's avatar
Patrick von Platen committed
431
432
433
        # 6. Add noise
        variance = 0
        if t > 0:
434
            device = model_output.device
435
436
437
            variance_noise = randn_tensor(
                model_output.shape, generator=generator, device=device, dtype=model_output.dtype
            )
438
439
            if self.variance_type == "fixed_small_log":
                variance = self._get_variance(t, predicted_variance=predicted_variance) * variance_noise
440
441
442
            elif self.variance_type == "learned_range":
                variance = self._get_variance(t, predicted_variance=predicted_variance)
                variance = torch.exp(0.5 * variance) * variance_noise
443
444
            else:
                variance = (self._get_variance(t, predicted_variance=predicted_variance) ** 0.5) * variance_noise
Patrick von Platen's avatar
Patrick von Platen committed
445
446
447

        pred_prev_sample = pred_prev_sample + variance

448
449
450
        if not return_dict:
            return (pred_prev_sample,)

451
        return DDPMSchedulerOutput(prev_sample=pred_prev_sample, pred_original_sample=pred_original_sample)
Patrick von Platen's avatar
Patrick von Platen committed
452

Partho's avatar
Partho committed
453
454
    def add_noise(
        self,
455
456
457
458
        original_samples: torch.FloatTensor,
        noise: torch.FloatTensor,
        timesteps: torch.IntTensor,
    ) -> torch.FloatTensor:
459
        # Make sure alphas_cumprod and timestep have same device and dtype as original_samples
460
        alphas_cumprod = self.alphas_cumprod.to(device=original_samples.device, dtype=original_samples.dtype)
461
        timesteps = timesteps.to(original_samples.device)
462

463
        sqrt_alpha_prod = alphas_cumprod[timesteps] ** 0.5
464
465
466
467
        sqrt_alpha_prod = sqrt_alpha_prod.flatten()
        while len(sqrt_alpha_prod.shape) < len(original_samples.shape):
            sqrt_alpha_prod = sqrt_alpha_prod.unsqueeze(-1)

468
        sqrt_one_minus_alpha_prod = (1 - alphas_cumprod[timesteps]) ** 0.5
469
        sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.flatten()
470
471
        while len(sqrt_one_minus_alpha_prod.shape) < len(original_samples.shape):
            sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.unsqueeze(-1)
472
473

        noisy_samples = sqrt_alpha_prod * original_samples + sqrt_one_minus_alpha_prod * noise
anton-l's avatar
anton-l committed
474
        return noisy_samples
anton-l's avatar
anton-l committed
475

476
477
478
479
    def get_velocity(
        self, sample: torch.FloatTensor, noise: torch.FloatTensor, timesteps: torch.IntTensor
    ) -> torch.FloatTensor:
        # Make sure alphas_cumprod and timestep have same device and dtype as sample
480
        alphas_cumprod = self.alphas_cumprod.to(device=sample.device, dtype=sample.dtype)
481
482
        timesteps = timesteps.to(sample.device)

483
        sqrt_alpha_prod = alphas_cumprod[timesteps] ** 0.5
484
485
486
487
        sqrt_alpha_prod = sqrt_alpha_prod.flatten()
        while len(sqrt_alpha_prod.shape) < len(sample.shape):
            sqrt_alpha_prod = sqrt_alpha_prod.unsqueeze(-1)

488
        sqrt_one_minus_alpha_prod = (1 - alphas_cumprod[timesteps]) ** 0.5
489
490
491
492
493
494
495
        sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.flatten()
        while len(sqrt_one_minus_alpha_prod.shape) < len(sample.shape):
            sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.unsqueeze(-1)

        velocity = sqrt_alpha_prod * noise - sqrt_one_minus_alpha_prod * sample
        return velocity

Patrick von Platen's avatar
improve  
Patrick von Platen committed
496
    def __len__(self):
Nathan Lambert's avatar
Nathan Lambert committed
497
        return self.config.num_train_timesteps
Will Berman's avatar
Will Berman committed
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512

    def previous_timestep(self, timestep):
        if self.custom_timesteps:
            index = (self.timesteps == timestep).nonzero(as_tuple=True)[0][0]
            if index == self.timesteps.shape[0] - 1:
                prev_t = torch.tensor(-1)
            else:
                prev_t = self.timesteps[index + 1]
        else:
            num_inference_steps = (
                self.num_inference_steps if self.num_inference_steps else self.config.num_train_timesteps
            )
            prev_t = timestep - self.config.num_train_timesteps // num_inference_steps

        return prev_t