|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
import math |
|
from typing import List, Optional, Tuple, Union |
|
|
|
import numpy as np |
|
import paddle |
|
|
|
from ..configuration_utils import ConfigMixin, register_to_config |
|
from ..utils import _COMPATIBLE_STABLE_DIFFUSION_SCHEDULERS, deprecate |
|
from .scheduling_utils import SchedulerMixin, SchedulerOutput |
|
|
|
|
|
def betas_for_alpha_bar(num_diffusion_timesteps, max_beta=0.999): |
|
""" |
|
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]. |
|
|
|
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 |
|
prevent singularities. |
|
|
|
Returns: |
|
betas (`np.ndarray`): the betas used by the scheduler to step the model outputs |
|
""" |
|
|
|
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 paddle.to_tensor(betas, dtype="float32") |
|
|
|
|
|
class DPMSolverMultistepScheduler(SchedulerMixin, ConfigMixin): |
|
""" |
|
DPM-Solver (and the improved version DPM-Solver++) is a fast dedicated high-order solver for diffusion ODEs with |
|
the convergence order guarantee. Empirically, sampling by DPM-Solver with only 20 steps can generate high-quality |
|
samples, and it can generate quite good samples even in only 10 steps. |
|
|
|
For more details, see the original paper: https://arxiv.org/abs/2206.00927 and https://arxiv.org/abs/2211.01095 |
|
|
|
Currently, we support the multistep DPM-Solver for both noise prediction models and data prediction models. We |
|
recommend to use `solver_order=2` for guided sampling, and `solver_order=3` for unconditional sampling. |
|
|
|
We also support the "dynamic thresholding" method in Imagen (https://arxiv.org/abs/2205.11487). For pixel-space |
|
diffusion models, you can set both `algorithm_type="dpmsolver++"` and `thresholding=True` to use the dynamic |
|
thresholding. Note that the thresholding method is unsuitable for latent-space diffusion models (such as |
|
stable-diffusion). |
|
|
|
[`~ConfigMixin`] takes care of storing all config attributes that are passed in the scheduler's `__init__` |
|
function, such as `num_train_timesteps`. They can be accessed via `scheduler.config.num_train_timesteps`. |
|
[`SchedulerMixin`] provides general loading and saving functionality via the [`SchedulerMixin.save_pretrained`] and |
|
[`~SchedulerMixin.from_pretrained`] functions. |
|
|
|
Args: |
|
num_train_timesteps (`int`): number of diffusion steps used to train the model. |
|
beta_start (`float`): the starting `beta` value of inference. |
|
beta_end (`float`): the final `beta` value. |
|
beta_schedule (`str`): |
|
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`. |
|
trained_betas (`np.ndarray`, optional): |
|
option to pass an array of betas directly to the constructor to bypass `beta_start`, `beta_end` etc. |
|
solver_order (`int`, default `2`): |
|
the order of DPM-Solver; can be `1` or `2` or `3`. We recommend to use `solver_order=2` for guided |
|
sampling, and `solver_order=3` for unconditional sampling. |
|
prediction_type (`str`, default `epsilon`, optional): |
|
prediction type of the scheduler function, one of `epsilon` (predicting the noise of the diffusion |
|
process), `sample` (directly predicting the noisy sample`) or `v_prediction` (see section 2.4 |
|
https://imagen.research.google/video/paper.pdf) |
|
thresholding (`bool`, default `False`): |
|
whether to use the "dynamic thresholding" method (introduced by Imagen, https://arxiv.org/abs/2205.11487). |
|
For pixel-space diffusion models, you can set both `algorithm_type=dpmsolver++` and `thresholding=True` to |
|
use the dynamic thresholding. Note that the thresholding method is unsuitable for latent-space diffusion |
|
models (such as stable-diffusion). |
|
dynamic_thresholding_ratio (`float`, default `0.995`): |
|
the ratio for the dynamic thresholding method. Default is `0.995`, the same as Imagen |
|
(https://arxiv.org/abs/2205.11487). |
|
sample_max_value (`float`, default `1.0`): |
|
the threshold value for dynamic thresholding. Valid only when `thresholding=True` and |
|
`algorithm_type="dpmsolver++`. |
|
algorithm_type (`str`, default `dpmsolver++`): |
|
the algorithm type for the solver. Either `dpmsolver` or `dpmsolver++`. The `dpmsolver` type implements the |
|
algorithms in https://arxiv.org/abs/2206.00927, and the `dpmsolver++` type implements the algorithms in |
|
https://arxiv.org/abs/2211.01095. We recommend to use `dpmsolver++` with `solver_order=2` for guided |
|
sampling (e.g. stable-diffusion). |
|
solver_type (`str`, default `midpoint`): |
|
the solver type for the second-order solver. Either `midpoint` or `heun`. The solver type slightly affects |
|
the sample quality, especially for small number of steps. We empirically find that `midpoint` solvers are |
|
slightly better, so we recommend to use the `midpoint` type. |
|
lower_order_final (`bool`, default `True`): |
|
whether to use lower-order solvers in the final steps. Only valid for < 15 inference steps. We empirically |
|
find this trick can stabilize the sampling of DPM-Solver for steps < 15, especially for steps <= 10. |
|
|
|
""" |
|
|
|
_compatibles = _COMPATIBLE_STABLE_DIFFUSION_SCHEDULERS.copy() |
|
_deprecated_kwargs = ["predict_epsilon"] |
|
order = 1 |
|
|
|
@register_to_config |
|
def __init__( |
|
self, |
|
num_train_timesteps: int = 1000, |
|
beta_start: float = 0.0001, |
|
beta_end: float = 0.02, |
|
beta_schedule: str = "linear", |
|
trained_betas: Optional[Union[np.ndarray, List[float]]] = None, |
|
solver_order: int = 2, |
|
prediction_type: str = "epsilon", |
|
thresholding: bool = False, |
|
dynamic_thresholding_ratio: float = 0.995, |
|
sample_max_value: float = 1.0, |
|
algorithm_type: str = "dpmsolver++", |
|
solver_type: str = "midpoint", |
|
lower_order_final: bool = True, |
|
**kwargs, |
|
): |
|
message = ( |
|
"Please make sure to instantiate your scheduler with `prediction_type` instead. E.g. `scheduler =" |
|
" DPMSolverMultistepScheduler.from_pretrained(<model_id>, prediction_type='epsilon')`." |
|
) |
|
predict_epsilon = deprecate("predict_epsilon", "0.13.0", message, take_from=kwargs) |
|
if predict_epsilon is not None: |
|
self.register_to_config(prediction_type="epsilon" if predict_epsilon else "sample") |
|
if trained_betas is not None: |
|
self.betas = paddle.to_tensor(trained_betas, dtype="float32") |
|
elif beta_schedule == "linear": |
|
self.betas = paddle.linspace(beta_start, beta_end, num_train_timesteps, dtype="float32") |
|
elif beta_schedule == "scaled_linear": |
|
|
|
self.betas = paddle.linspace(beta_start**0.5, beta_end**0.5, num_train_timesteps, dtype="float32") ** 2 |
|
elif beta_schedule == "squaredcos_cap_v2": |
|
|
|
self.betas = betas_for_alpha_bar(num_train_timesteps) |
|
else: |
|
raise NotImplementedError(f"{beta_schedule} does is not implemented for {self.__class__}") |
|
|
|
self.alphas = 1.0 - self.betas |
|
self.alphas_cumprod = paddle.cumprod(self.alphas, 0) |
|
|
|
self.alpha_t = paddle.sqrt(self.alphas_cumprod) |
|
self.sigma_t = paddle.sqrt(1 - self.alphas_cumprod) |
|
self.lambda_t = paddle.log(self.alpha_t) - paddle.log(self.sigma_t) |
|
|
|
|
|
self.init_noise_sigma = 1.0 |
|
|
|
|
|
if algorithm_type not in ["dpmsolver", "dpmsolver++"]: |
|
raise NotImplementedError(f"{algorithm_type} does is not implemented for {self.__class__}") |
|
if solver_type not in ["midpoint", "heun"]: |
|
raise NotImplementedError(f"{solver_type} does is not implemented for {self.__class__}") |
|
|
|
|
|
self.num_inference_steps = None |
|
timesteps = np.linspace(0, num_train_timesteps - 1, num_train_timesteps, dtype=np.float32)[::-1].copy() |
|
self.timesteps = paddle.to_tensor(timesteps) |
|
self.model_outputs = [None] * solver_order |
|
self.lower_order_nums = 0 |
|
|
|
def set_timesteps(self, num_inference_steps: int): |
|
""" |
|
Sets the timesteps used for the diffusion chain. Supporting function to be run before inference. |
|
|
|
Args: |
|
num_inference_steps (`int`): |
|
the number of diffusion steps used when generating samples with a pre-trained model. |
|
""" |
|
self.num_inference_steps = num_inference_steps |
|
timesteps = ( |
|
np.linspace(0, self.num_train_timesteps - 1, num_inference_steps + 1) |
|
.round()[::-1][:-1] |
|
.copy() |
|
.astype(np.int64) |
|
) |
|
self.timesteps = paddle.to_tensor(timesteps) |
|
self.model_outputs = [ |
|
None, |
|
] * self.config.solver_order |
|
self.lower_order_nums = 0 |
|
|
|
def convert_model_output(self, model_output: paddle.Tensor, timestep: int, sample: paddle.Tensor) -> paddle.Tensor: |
|
""" |
|
Convert the model output to the corresponding type that the algorithm (DPM-Solver / DPM-Solver++) needs. |
|
|
|
DPM-Solver is designed to discretize an integral of the noise prediction model, and DPM-Solver++ is designed to |
|
discretize an integral of the data prediction model. So we need to first convert the model output to the |
|
corresponding type to match the algorithm. |
|
|
|
Note that the algorithm type and the model type is decoupled. That is to say, we can use either DPM-Solver or |
|
DPM-Solver++ for both noise prediction model and data prediction model. |
|
|
|
Args: |
|
model_output (`paddle.Tensor`): direct output from learned diffusion model. |
|
timestep (`int`): current discrete timestep in the diffusion chain. |
|
sample (`paddle.Tensor`): |
|
current instance of sample being created by diffusion process. |
|
|
|
Returns: |
|
`paddle.Tensor`: the converted model output. |
|
""" |
|
|
|
if self.config.algorithm_type == "dpmsolver++": |
|
if self.config.prediction_type == "epsilon": |
|
alpha_t, sigma_t = self.alpha_t[timestep], self.sigma_t[timestep] |
|
x0_pred = (sample - sigma_t * model_output) / alpha_t |
|
elif self.config.prediction_type == "sample": |
|
x0_pred = model_output |
|
elif self.config.prediction_type == "v_prediction": |
|
alpha_t, sigma_t = self.alpha_t[timestep], self.sigma_t[timestep] |
|
x0_pred = alpha_t * sample - sigma_t * model_output |
|
else: |
|
raise ValueError( |
|
f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`, or" |
|
" `v_prediction` for the DPMSolverMultistepScheduler." |
|
) |
|
|
|
if self.config.thresholding: |
|
|
|
orig_dtype = x0_pred.dtype |
|
if orig_dtype not in [paddle.float32, paddle.float64]: |
|
x0_pred = x0_pred.cast("float32") |
|
dynamic_max_val = paddle.quantile( |
|
paddle.abs(x0_pred).reshape((x0_pred.shape[0], -1)), self.config.dynamic_thresholding_ratio, axis=1 |
|
) |
|
dynamic_max_val = paddle.maximum( |
|
dynamic_max_val, |
|
self.config.sample_max_value * paddle.ones_like(dynamic_max_val), |
|
)[(...,) + (None,) * (x0_pred.ndim - 1)] |
|
x0_pred = paddle.clip(x0_pred, -dynamic_max_val, dynamic_max_val) / dynamic_max_val |
|
x0_pred = x0_pred.cast(orig_dtype) |
|
return x0_pred |
|
|
|
elif self.config.algorithm_type == "dpmsolver": |
|
if self.config.prediction_type == "epsilon": |
|
return model_output |
|
elif self.config.prediction_type == "sample": |
|
alpha_t, sigma_t = self.alpha_t[timestep], self.sigma_t[timestep] |
|
epsilon = (sample - alpha_t * model_output) / sigma_t |
|
return epsilon |
|
elif self.config.prediction_type == "v_prediction": |
|
alpha_t, sigma_t = self.alpha_t[timestep], self.sigma_t[timestep] |
|
epsilon = alpha_t * model_output + sigma_t * sample |
|
return epsilon |
|
else: |
|
raise ValueError( |
|
f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`, or" |
|
" `v_prediction` for the DPMSolverMultistepScheduler." |
|
) |
|
|
|
def dpm_solver_first_order_update( |
|
self, |
|
model_output: paddle.Tensor, |
|
timestep: int, |
|
prev_timestep: int, |
|
sample: paddle.Tensor, |
|
) -> paddle.Tensor: |
|
""" |
|
One step for the first-order DPM-Solver (equivalent to DDIM). |
|
|
|
See https://arxiv.org/abs/2206.00927 for the detailed derivation. |
|
|
|
Args: |
|
model_output (`paddle.Tensor`): direct output from learned diffusion model. |
|
timestep (`int`): current discrete timestep in the diffusion chain. |
|
prev_timestep (`int`): previous discrete timestep in the diffusion chain. |
|
sample (`paddle.Tensor`): |
|
current instance of sample being created by diffusion process. |
|
|
|
Returns: |
|
`paddle.Tensor`: the sample tensor at the previous timestep. |
|
""" |
|
lambda_t, lambda_s = self.lambda_t[prev_timestep], self.lambda_t[timestep] |
|
alpha_t, alpha_s = self.alpha_t[prev_timestep], self.alpha_t[timestep] |
|
sigma_t, sigma_s = self.sigma_t[prev_timestep], self.sigma_t[timestep] |
|
h = lambda_t - lambda_s |
|
if self.config.algorithm_type == "dpmsolver++": |
|
x_t = (sigma_t / sigma_s) * sample - (alpha_t * (paddle.exp(-h) - 1.0)) * model_output |
|
elif self.config.algorithm_type == "dpmsolver": |
|
x_t = (alpha_t / alpha_s) * sample - (sigma_t * (paddle.exp(h) - 1.0)) * model_output |
|
return x_t |
|
|
|
def multistep_dpm_solver_second_order_update( |
|
self, |
|
model_output_list: List[paddle.Tensor], |
|
timestep_list: List[int], |
|
prev_timestep: int, |
|
sample: paddle.Tensor, |
|
) -> paddle.Tensor: |
|
""" |
|
One step for the second-order multistep DPM-Solver. |
|
|
|
Args: |
|
model_output_list (`List[paddle.Tensor]`): |
|
direct outputs from learned diffusion model at current and latter timesteps. |
|
timestep (`int`): current and latter discrete timestep in the diffusion chain. |
|
prev_timestep (`int`): previous discrete timestep in the diffusion chain. |
|
sample (`paddle.Tensor`): |
|
current instance of sample being created by diffusion process. |
|
|
|
Returns: |
|
`paddle.Tensor`: the sample tensor at the previous timestep. |
|
""" |
|
t, s0, s1 = prev_timestep, timestep_list[-1], timestep_list[-2] |
|
m0, m1 = model_output_list[-1], model_output_list[-2] |
|
lambda_t, lambda_s0, lambda_s1 = self.lambda_t[t], self.lambda_t[s0], self.lambda_t[s1] |
|
alpha_t, alpha_s0 = self.alpha_t[t], self.alpha_t[s0] |
|
sigma_t, sigma_s0 = self.sigma_t[t], self.sigma_t[s0] |
|
h, h_0 = lambda_t - lambda_s0, lambda_s0 - lambda_s1 |
|
r0 = h_0 / h |
|
D0, D1 = m0, (1.0 / r0) * (m0 - m1) |
|
if self.config.algorithm_type == "dpmsolver++": |
|
|
|
if self.config.solver_type == "midpoint": |
|
x_t = ( |
|
(sigma_t / sigma_s0) * sample |
|
- (alpha_t * (paddle.exp(-h) - 1.0)) * D0 |
|
- 0.5 * (alpha_t * (paddle.exp(-h) - 1.0)) * D1 |
|
) |
|
elif self.config.solver_type == "heun": |
|
x_t = ( |
|
(sigma_t / sigma_s0) * sample |
|
- (alpha_t * (paddle.exp(-h) - 1.0)) * D0 |
|
+ (alpha_t * ((paddle.exp(-h) - 1.0) / h + 1.0)) * D1 |
|
) |
|
elif self.config.algorithm_type == "dpmsolver": |
|
|
|
if self.config.solver_type == "midpoint": |
|
x_t = ( |
|
(alpha_t / alpha_s0) * sample |
|
- (sigma_t * (paddle.exp(h) - 1.0)) * D0 |
|
- 0.5 * (sigma_t * (paddle.exp(h) - 1.0)) * D1 |
|
) |
|
elif self.config.solver_type == "heun": |
|
x_t = ( |
|
(alpha_t / alpha_s0) * sample |
|
- (sigma_t * (paddle.exp(h) - 1.0)) * D0 |
|
- (sigma_t * ((paddle.exp(h) - 1.0) / h - 1.0)) * D1 |
|
) |
|
return x_t |
|
|
|
def multistep_dpm_solver_third_order_update( |
|
self, |
|
model_output_list: List[paddle.Tensor], |
|
timestep_list: List[int], |
|
prev_timestep: int, |
|
sample: paddle.Tensor, |
|
) -> paddle.Tensor: |
|
""" |
|
One step for the third-order multistep DPM-Solver. |
|
|
|
Args: |
|
model_output_list (`List[paddle.Tensor]`): |
|
direct outputs from learned diffusion model at current and latter timesteps. |
|
timestep (`int`): current and latter discrete timestep in the diffusion chain. |
|
prev_timestep (`int`): previous discrete timestep in the diffusion chain. |
|
sample (`paddle.Tensor`): |
|
current instance of sample being created by diffusion process. |
|
|
|
Returns: |
|
`paddle.Tensor`: the sample tensor at the previous timestep. |
|
""" |
|
t, s0, s1, s2 = prev_timestep, timestep_list[-1], timestep_list[-2], timestep_list[-3] |
|
m0, m1, m2 = model_output_list[-1], model_output_list[-2], model_output_list[-3] |
|
lambda_t, lambda_s0, lambda_s1, lambda_s2 = ( |
|
self.lambda_t[t], |
|
self.lambda_t[s0], |
|
self.lambda_t[s1], |
|
self.lambda_t[s2], |
|
) |
|
alpha_t, alpha_s0 = self.alpha_t[t], self.alpha_t[s0] |
|
sigma_t, sigma_s0 = self.sigma_t[t], self.sigma_t[s0] |
|
h, h_0, h_1 = lambda_t - lambda_s0, lambda_s0 - lambda_s1, lambda_s1 - lambda_s2 |
|
r0, r1 = h_0 / h, h_1 / h |
|
D0 = m0 |
|
D1_0, D1_1 = (1.0 / r0) * (m0 - m1), (1.0 / r1) * (m1 - m2) |
|
D1 = D1_0 + (r0 / (r0 + r1)) * (D1_0 - D1_1) |
|
D2 = (1.0 / (r0 + r1)) * (D1_0 - D1_1) |
|
if self.config.algorithm_type == "dpmsolver++": |
|
|
|
x_t = ( |
|
(sigma_t / sigma_s0) * sample |
|
- (alpha_t * (paddle.exp(-h) - 1.0)) * D0 |
|
+ (alpha_t * ((paddle.exp(-h) - 1.0) / h + 1.0)) * D1 |
|
- (alpha_t * ((paddle.exp(-h) - 1.0 + h) / h**2 - 0.5)) * D2 |
|
) |
|
elif self.config.algorithm_type == "dpmsolver": |
|
|
|
x_t = ( |
|
(alpha_t / alpha_s0) * sample |
|
- (sigma_t * (paddle.exp(h) - 1.0)) * D0 |
|
- (sigma_t * ((paddle.exp(h) - 1.0) / h - 1.0)) * D1 |
|
- (sigma_t * ((paddle.exp(h) - 1.0 - h) / h**2 - 0.5)) * D2 |
|
) |
|
return x_t |
|
|
|
def step( |
|
self, |
|
model_output: paddle.Tensor, |
|
timestep: int, |
|
sample: paddle.Tensor, |
|
return_dict: bool = True, |
|
) -> Union[SchedulerOutput, Tuple]: |
|
""" |
|
Step function propagating the sample with the multistep DPM-Solver. |
|
|
|
Args: |
|
model_output (`paddle.Tensor`): direct output from learned diffusion model. |
|
timestep (`int`): current discrete timestep in the diffusion chain. |
|
sample (`paddle.Tensor`): |
|
current instance of sample being created by diffusion process. |
|
return_dict (`bool`): option for returning tuple rather than SchedulerOutput class |
|
|
|
Returns: |
|
[`~scheduling_utils.SchedulerOutput`] or `tuple`: [`~scheduling_utils.SchedulerOutput`] if `return_dict` is |
|
True, otherwise a `tuple`. When returning a tuple, the first element is the sample tensor. |
|
|
|
""" |
|
if self.num_inference_steps is None: |
|
raise ValueError( |
|
"Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler" |
|
) |
|
|
|
step_index = (self.timesteps == timestep).nonzero() |
|
if len(step_index) == 0: |
|
step_index = len(self.timesteps) - 1 |
|
else: |
|
step_index = step_index.item() |
|
prev_timestep = 0 if step_index == len(self.timesteps) - 1 else self.timesteps[step_index + 1] |
|
lower_order_final = ( |
|
(step_index == len(self.timesteps) - 1) and self.config.lower_order_final and len(self.timesteps) < 15 |
|
) |
|
lower_order_second = ( |
|
(step_index == len(self.timesteps) - 2) and self.config.lower_order_final and len(self.timesteps) < 15 |
|
) |
|
|
|
model_output = self.convert_model_output(model_output, timestep, sample) |
|
for i in range(self.config.solver_order - 1): |
|
self.model_outputs[i] = self.model_outputs[i + 1] |
|
self.model_outputs[-1] = model_output |
|
|
|
if self.config.solver_order == 1 or self.lower_order_nums < 1 or lower_order_final: |
|
prev_sample = self.dpm_solver_first_order_update(model_output, timestep, prev_timestep, sample) |
|
elif self.config.solver_order == 2 or self.lower_order_nums < 2 or lower_order_second: |
|
timestep_list = [self.timesteps[step_index - 1], timestep] |
|
prev_sample = self.multistep_dpm_solver_second_order_update( |
|
self.model_outputs, timestep_list, prev_timestep, sample |
|
) |
|
else: |
|
timestep_list = [self.timesteps[step_index - 2], self.timesteps[step_index - 1], timestep] |
|
prev_sample = self.multistep_dpm_solver_third_order_update( |
|
self.model_outputs, timestep_list, prev_timestep, sample |
|
) |
|
|
|
if self.lower_order_nums < self.config.solver_order: |
|
self.lower_order_nums += 1 |
|
|
|
if not return_dict: |
|
return (prev_sample,) |
|
|
|
return SchedulerOutput(prev_sample=prev_sample) |
|
|
|
def scale_model_input(self, sample: paddle.Tensor, *args, **kwargs) -> paddle.Tensor: |
|
""" |
|
Ensures interchangeability with schedulers that need to scale the denoising model input depending on the |
|
current timestep. |
|
|
|
Args: |
|
sample (`paddle.Tensor`): input sample |
|
|
|
Returns: |
|
`paddle.Tensor`: scaled input sample |
|
""" |
|
return sample |
|
|
|
def add_noise( |
|
self, |
|
original_samples: paddle.Tensor, |
|
noise: paddle.Tensor, |
|
timesteps: paddle.Tensor, |
|
) -> paddle.Tensor: |
|
|
|
self.alphas_cumprod = self.alphas_cumprod.cast(original_samples.dtype) |
|
|
|
sqrt_alpha_prod = self.alphas_cumprod[timesteps] ** 0.5 |
|
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) |
|
|
|
sqrt_one_minus_alpha_prod = (1 - self.alphas_cumprod[timesteps]) ** 0.5 |
|
sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.flatten() |
|
while len(sqrt_one_minus_alpha_prod.shape) < len(original_samples.shape): |
|
sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.unsqueeze(-1) |
|
|
|
noisy_samples = sqrt_alpha_prod * original_samples + sqrt_one_minus_alpha_prod * noise |
|
return noisy_samples |
|
|
|
def __len__(self): |
|
return self.config.num_train_timesteps |
|
|