scheduling_ddpm.py 25.4 KB
Newer Older
Aryan's avatar
Aryan committed
1
# Copyright 2025 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

    Args:
36
        prev_sample (`torch.Tensor` 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
            denoising loop.
39
        pred_original_sample (`torch.Tensor` 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
            `pred_original_sample` can be used to preview progress or for guidance.
    """

44
45
    prev_sample: torch.Tensor
    pred_original_sample: Optional[torch.Tensor] = 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
        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:
M. Tolga Cangöz's avatar
M. Tolga Cangöz committed
82
        raise ValueError(f"Unsupported alpha_transform_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


92
93
94
# Copied from diffusers.schedulers.scheduling_ddim.rescale_zero_terminal_snr
def rescale_zero_terminal_snr(betas):
    """
Quentin Gallouédec's avatar
Quentin Gallouédec committed
95
    Rescales betas to have zero terminal SNR Based on https://huggingface.co/papers/2305.08891 (Algorithm 1)
96
97

    Args:
98
        betas (`torch.Tensor`):
99
100
101
            the betas that the scheduler is being initialized with.

    Returns:
102
        `torch.Tensor`: rescaled betas with zero terminal SNR
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
    """
    # Convert betas to alphas_bar_sqrt
    alphas = 1.0 - betas
    alphas_cumprod = torch.cumprod(alphas, dim=0)
    alphas_bar_sqrt = alphas_cumprod.sqrt()

    # Store old values.
    alphas_bar_sqrt_0 = alphas_bar_sqrt[0].clone()
    alphas_bar_sqrt_T = alphas_bar_sqrt[-1].clone()

    # Shift so the last timestep is zero.
    alphas_bar_sqrt -= alphas_bar_sqrt_T

    # Scale so the first timestep is back to the old value.
    alphas_bar_sqrt *= alphas_bar_sqrt_0 / (alphas_bar_sqrt_0 - alphas_bar_sqrt_T)

    # Convert alphas_bar_sqrt to betas
    alphas_bar = alphas_bar_sqrt**2  # Revert sqrt
    alphas = alphas_bar[1:] / alphas_bar[:-1]  # Revert cumprod
    alphas = torch.cat([alphas_bar[0:1], alphas])
    betas = 1 - alphas

    return betas


Patrick von Platen's avatar
Patrick von Platen committed
128
class DDPMScheduler(SchedulerMixin, ConfigMixin):
129
    """
130
    `DDPMScheduler` explores the connections between denoising score matching and Langevin dynamics sampling.
131

132
133
    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.
134
135

    Args:
136
137
138
139
140
141
142
143
        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
144
            `linear`, `scaled_linear`, `squaredcos_cap_v2`, or `sigmoid`.
145
146
        trained_betas (`np.ndarray`, *optional*):
            An array of betas to pass directly to the constructor without using `beta_start` and `beta_end`.
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
        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):
169
            An offset added to the inference steps, as required by some model families.
170
171
172
173
        rescale_betas_zero_snr (`bool`, defaults to `False`):
            Whether to rescale the betas to have zero terminal SNR. This enables the model to generate very bright and
            dark samples instead of limiting it to samples with medium brightness. Loosely related to
            [`--offset_noise`](https://github.com/huggingface/diffusers/blob/74fd735eb073eb1d774b1ab4154a0876eb82f055/examples/dreambooth/train_dreambooth.py#L506).
174
175
    """

Kashif Rasul's avatar
Kashif Rasul committed
176
    _compatibles = [e.name for e in KarrasDiffusionSchedulers]
177
    order = 1
178

179
    @register_to_config
Patrick von Platen's avatar
improve  
Patrick von Platen committed
180
181
    def __init__(
        self,
Partho's avatar
Partho committed
182
183
184
185
        num_train_timesteps: int = 1000,
        beta_start: float = 0.0001,
        beta_end: float = 0.02,
        beta_schedule: str = "linear",
186
        trained_betas: Optional[Union[np.ndarray, List[float]]] = None,
Partho's avatar
Partho committed
187
188
        variance_type: str = "fixed_small",
        clip_sample: bool = True,
189
        prediction_type: str = "epsilon",
190
191
192
193
        thresholding: bool = False,
        dynamic_thresholding_ratio: float = 0.995,
        clip_sample_range: float = 1.0,
        sample_max_value: float = 1.0,
194
195
        timestep_spacing: str = "leading",
        steps_offset: int = 0,
196
        rescale_betas_zero_snr: bool = False,
Patrick von Platen's avatar
improve  
Patrick von Platen committed
197
    ):
198
        if trained_betas is not None:
199
            self.betas = torch.tensor(trained_betas, dtype=torch.float32)
200
        elif beta_schedule == "linear":
201
            self.betas = torch.linspace(beta_start, beta_end, num_train_timesteps, dtype=torch.float32)
202
203
        elif beta_schedule == "scaled_linear":
            # this schedule is very specific to the latent diffusion model.
204
            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
205
        elif beta_schedule == "squaredcos_cap_v2":
Patrick von Platen's avatar
Patrick von Platen committed
206
            # Glide cosine schedule
Nathan Lambert's avatar
Nathan Lambert committed
207
            self.betas = betas_for_alpha_bar(num_train_timesteps)
Nathan Lambert's avatar
Nathan Lambert committed
208
209
210
211
        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
212
        else:
213
            raise NotImplementedError(f"{beta_schedule} is not implemented for {self.__class__}")
Patrick von Platen's avatar
improve  
Patrick von Platen committed
214

215
216
217
218
        # Rescale for zero SNR
        if rescale_betas_zero_snr:
            self.betas = rescale_zero_terminal_snr(self.betas)

Patrick von Platen's avatar
Patrick von Platen committed
219
        self.alphas = 1.0 - self.betas
220
221
        self.alphas_cumprod = torch.cumprod(self.alphas, dim=0)
        self.one = torch.tensor(1.0)
Patrick von Platen's avatar
Patrick von Platen committed
222

223
224
225
        # standard deviation of the initial noise distribution
        self.init_noise_sigma = 1.0

226
        # setable values
Will Berman's avatar
Will Berman committed
227
        self.custom_timesteps = False
228
        self.num_inference_steps = None
229
        self.timesteps = torch.from_numpy(np.arange(0, num_train_timesteps)[::-1].copy())
Patrick von Platen's avatar
Patrick von Platen committed
230

231
232
        self.variance_type = variance_type

233
    def scale_model_input(self, sample: torch.Tensor, timestep: Optional[int] = None) -> torch.Tensor:
234
235
236
237
238
        """
        Ensures interchangeability with schedulers that need to scale the denoising model input depending on the
        current timestep.

        Args:
239
            sample (`torch.Tensor`):
240
241
242
                The input sample.
            timestep (`int`, *optional*):
                The current timestep in the diffusion chain.
243
244

        Returns:
245
            `torch.Tensor`:
246
                A scaled input sample.
247
248
249
        """
        return sample

Will Berman's avatar
Will Berman committed
250
251
252
253
254
255
    def set_timesteps(
        self,
        num_inference_steps: Optional[int] = None,
        device: Union[str, torch.device] = None,
        timesteps: Optional[List[int]] = None,
    ):
256
        """
257
        Sets the discrete timesteps used for the diffusion chain (to be run before inference).
258
259

        Args:
260
261
            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
262
                `timesteps` must be `None`.
263
264
265
266
267
268
            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
269

270
        """
Will Berman's avatar
Will Berman committed
271
272
273
274
275
276
277
278
279
280
        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(
281
                    f"`timesteps` must start before `self.config.train_timesteps`: {self.config.num_train_timesteps}."
Will Berman's avatar
Will Berman committed
282
283
284
285
286
287
288
289
290
291
292
                )

            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."
                )
293

Will Berman's avatar
Will Berman committed
294
295
            self.num_inference_steps = num_inference_steps
            self.custom_timesteps = False
296

Quentin Gallouédec's avatar
Quentin Gallouédec committed
297
            # "linspace", "leading", "trailing" corresponds to annotation of Table 2. of https://huggingface.co/papers/2305.08891
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
            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'."
                )

322
        self.timesteps = torch.from_numpy(timesteps).to(device)
323

324
    def _get_variance(self, t, predicted_variance=None, variance_type=None):
Will Berman's avatar
Will Berman committed
325
326
        prev_t = self.previous_timestep(t)

327
        alpha_prod_t = self.alphas_cumprod[t]
328
329
        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
330

Quentin Gallouédec's avatar
Quentin Gallouédec committed
331
        # For t > 0, compute predicted variance βt (see formula (6) and (7) from https://huggingface.co/papers/2006.11239)
332
        # and sample from it to get previous sample
Kashif Rasul's avatar
Kashif Rasul committed
333
        # x_{t-1} ~ N(pred_prev_sample, variance) == add variance to pred_sample
334
        variance = (1 - alpha_prod_t_prev) / (1 - alpha_prod_t) * current_beta_t
William Berman's avatar
William Berman committed
335
336

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

339
340
341
        if variance_type is None:
            variance_type = self.config.variance_type

342
        # hacks - were probably added for training stability
343
        if variance_type == "fixed_small":
William Berman's avatar
William Berman committed
344
            variance = variance
Quentin Gallouédec's avatar
Quentin Gallouédec committed
345
        # for rl-diffuser https://huggingface.co/papers/2205.09991
346
        elif variance_type == "fixed_small_log":
347
            variance = torch.log(variance)
348
            variance = torch.exp(0.5 * variance)
349
        elif variance_type == "fixed_large":
350
            variance = current_beta_t
351
        elif variance_type == "fixed_large_log":
Patrick von Platen's avatar
Patrick von Platen committed
352
            # Glide max_log
353
            variance = torch.log(current_beta_t)
354
355
356
        elif variance_type == "learned":
            return predicted_variance
        elif variance_type == "learned_range":
357
            min_log = torch.log(variance)
William Berman's avatar
William Berman committed
358
            max_log = torch.log(current_beta_t)
359
360
            frac = (predicted_variance + 1) / 2
            variance = frac * max_log + (1 - frac) * min_log
Patrick von Platen's avatar
Patrick von Platen committed
361
362
363

        return variance

364
    def _threshold_sample(self, sample: torch.Tensor) -> torch.Tensor:
365
366
367
368
369
370
371
        """
        "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."

Quentin Gallouédec's avatar
Quentin Gallouédec committed
372
        https://huggingface.co/papers/2205.11487
373
374
        """
        dtype = sample.dtype
375
        batch_size, channels, *remaining_dims = sample.shape
376
377
378
379
380

        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
381
        sample = sample.reshape(batch_size, channels * np.prod(remaining_dims))
382
383
384
385
386
387
388
389
390
391

        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"

392
        sample = sample.reshape(batch_size, channels, *remaining_dims)
393
394
395
        sample = sample.to(dtype)

        return sample
396

397
398
    def step(
        self,
399
        model_output: torch.Tensor,
400
        timestep: int,
401
        sample: torch.Tensor,
Patrick von Platen's avatar
Patrick von Platen committed
402
        generator=None,
403
        return_dict: bool = True,
404
    ) -> Union[DDPMSchedulerOutput, Tuple]:
405
        """
406
        Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion
407
408
409
        process from the learned model outputs (most often the predicted noise).

        Args:
410
            model_output (`torch.Tensor`):
411
412
413
                The direct output from learned diffusion model.
            timestep (`float`):
                The current discrete timestep in the diffusion chain.
414
            sample (`torch.Tensor`):
415
416
417
418
419
                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`.
420
421

        Returns:
422
423
424
            [`~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.
425
426

        """
427
        t = timestep
Will Berman's avatar
Will Berman committed
428
429

        prev_t = self.previous_timestep(t)
430

431
432
433
434
435
        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
436
        # 1. compute alphas, betas
437
        alpha_prod_t = self.alphas_cumprod[t]
438
        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
439
440
        beta_prod_t = 1 - alpha_prod_t
        beta_prod_t_prev = 1 - alpha_prod_t_prev
441
442
        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
443

444
        # 2. compute predicted original sample from predicted noise also called
Quentin Gallouédec's avatar
Quentin Gallouédec committed
445
        # "predicted x_0" of formula (15) from https://huggingface.co/papers/2006.11239
446
        if self.config.prediction_type == "epsilon":
Patrick von Platen's avatar
Patrick von Platen committed
447
            pred_original_sample = (sample - beta_prod_t ** (0.5) * model_output) / alpha_prod_t ** (0.5)
448
        elif self.config.prediction_type == "sample":
Patrick von Platen's avatar
Patrick von Platen committed
449
            pred_original_sample = model_output
450
451
        elif self.config.prediction_type == "v_prediction":
            pred_original_sample = (alpha_prod_t**0.5) * sample - (beta_prod_t**0.5) * model_output
452
453
        else:
            raise ValueError(
454
455
                f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample` or"
                " `v_prediction`  for the DDPMScheduler."
456
            )
Patrick von Platen's avatar
Patrick von Platen committed
457

458
        # 3. Clip or threshold "predicted x_0"
459
460
461
        if self.config.thresholding:
            pred_original_sample = self._threshold_sample(pred_original_sample)
        elif self.config.clip_sample:
462
463
            pred_original_sample = pred_original_sample.clamp(
                -self.config.clip_sample_range, self.config.clip_sample_range
Will Berman's avatar
Will Berman committed
464
            )
Patrick von Platen's avatar
Patrick von Platen committed
465

466
        # 4. Compute coefficients for pred_original_sample x_0 and current sample x_t
Quentin Gallouédec's avatar
Quentin Gallouédec committed
467
        # See formula (7) from https://huggingface.co/papers/2006.11239
468
469
        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
470

471
        # 5. Compute predicted previous sample µ_t
Quentin Gallouédec's avatar
Quentin Gallouédec committed
472
        # See formula (7) from https://huggingface.co/papers/2006.11239
473
        pred_prev_sample = pred_original_sample_coeff * pred_original_sample + current_sample_coeff * sample
Patrick von Platen's avatar
Patrick von Platen committed
474

Patrick von Platen's avatar
Patrick von Platen committed
475
476
477
        # 6. Add noise
        variance = 0
        if t > 0:
478
            device = model_output.device
479
480
481
            variance_noise = randn_tensor(
                model_output.shape, generator=generator, device=device, dtype=model_output.dtype
            )
482
483
            if self.variance_type == "fixed_small_log":
                variance = self._get_variance(t, predicted_variance=predicted_variance) * variance_noise
484
485
486
            elif self.variance_type == "learned_range":
                variance = self._get_variance(t, predicted_variance=predicted_variance)
                variance = torch.exp(0.5 * variance) * variance_noise
487
488
            else:
                variance = (self._get_variance(t, predicted_variance=predicted_variance) ** 0.5) * variance_noise
Patrick von Platen's avatar
Patrick von Platen committed
489
490
491

        pred_prev_sample = pred_prev_sample + variance

492
        if not return_dict:
493
494
495
496
            return (
                pred_prev_sample,
                pred_original_sample,
            )
497

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

Partho's avatar
Partho committed
500
501
    def add_noise(
        self,
502
503
        original_samples: torch.Tensor,
        noise: torch.Tensor,
504
        timesteps: torch.IntTensor,
505
    ) -> torch.Tensor:
506
        # Make sure alphas_cumprod and timestep have same device and dtype as original_samples
507
508
509
510
        # Move the self.alphas_cumprod to device to avoid redundant CPU to GPU data movement
        # for the subsequent add_noise calls
        self.alphas_cumprod = self.alphas_cumprod.to(device=original_samples.device)
        alphas_cumprod = self.alphas_cumprod.to(dtype=original_samples.dtype)
511
        timesteps = timesteps.to(original_samples.device)
512

513
        sqrt_alpha_prod = alphas_cumprod[timesteps] ** 0.5
514
515
516
517
        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)

518
        sqrt_one_minus_alpha_prod = (1 - alphas_cumprod[timesteps]) ** 0.5
519
        sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.flatten()
520
521
        while len(sqrt_one_minus_alpha_prod.shape) < len(original_samples.shape):
            sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.unsqueeze(-1)
522
523

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

526
    def get_velocity(self, sample: torch.Tensor, noise: torch.Tensor, timesteps: torch.IntTensor) -> torch.Tensor:
527
        # Make sure alphas_cumprod and timestep have same device and dtype as sample
528
529
        self.alphas_cumprod = self.alphas_cumprod.to(device=sample.device)
        alphas_cumprod = self.alphas_cumprod.to(dtype=sample.dtype)
530
531
        timesteps = timesteps.to(sample.device)

532
        sqrt_alpha_prod = alphas_cumprod[timesteps] ** 0.5
533
534
535
536
        sqrt_alpha_prod = sqrt_alpha_prod.flatten()
        while len(sqrt_alpha_prod.shape) < len(sample.shape):
            sqrt_alpha_prod = sqrt_alpha_prod.unsqueeze(-1)

537
        sqrt_one_minus_alpha_prod = (1 - alphas_cumprod[timesteps]) ** 0.5
538
539
540
541
542
543
544
        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
545
    def __len__(self):
Nathan Lambert's avatar
Nathan Lambert committed
546
        return self.config.num_train_timesteps
Will Berman's avatar
Will Berman committed
547
548

    def previous_timestep(self, timestep):
549
        if self.custom_timesteps or self.num_inference_steps:
Will Berman's avatar
Will Berman committed
550
551
552
553
554
555
            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:
556
            prev_t = timestep - 1
Will Berman's avatar
Will Berman committed
557
        return prev_t