|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
import math |
|
from dataclasses import dataclass |
|
from typing import List, Optional, Tuple, Union |
|
|
|
import numpy as np |
|
import torch |
|
|
|
from diffusers.configuration_utils import ConfigMixin, register_to_config |
|
from diffusers.utils import BaseOutput, logging |
|
from diffusers.utils.torch_utils import randn_tensor |
|
from diffusers.schedulers.scheduling_utils import SchedulerMixin |
|
|
|
|
|
logger = logging.get_logger(__name__) |
|
|
|
|
|
@dataclass |
|
class LCMSingleStepSchedulerOutput(BaseOutput): |
|
""" |
|
Output class for the scheduler's `step` function output. |
|
|
|
Args: |
|
pred_original_sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` for images): |
|
The predicted denoised sample `(x_{0})` based on the model output from the current timestep. |
|
`pred_original_sample` can be used to preview progress or for guidance. |
|
""" |
|
|
|
denoised: Optional[torch.FloatTensor] = None |
|
|
|
|
|
|
|
def betas_for_alpha_bar( |
|
num_diffusion_timesteps, |
|
max_beta=0.999, |
|
alpha_transform_type="cosine", |
|
): |
|
""" |
|
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. |
|
alpha_transform_type (`str`, *optional*, default to `cosine`): the type of noise schedule for alpha_bar. |
|
Choose from `cosine` or `exp` |
|
|
|
Returns: |
|
betas (`np.ndarray`): the betas used by the scheduler to step the model outputs |
|
""" |
|
if alpha_transform_type == "cosine": |
|
|
|
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}") |
|
|
|
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_fn(t2) / alpha_bar_fn(t1), max_beta)) |
|
return torch.tensor(betas, dtype=torch.float32) |
|
|
|
|
|
|
|
def rescale_zero_terminal_snr(betas: torch.FloatTensor) -> torch.FloatTensor: |
|
""" |
|
Rescales betas to have zero terminal SNR Based on https://arxiv.org/pdf/2305.08891.pdf (Algorithm 1) |
|
|
|
|
|
Args: |
|
betas (`torch.FloatTensor`): |
|
the betas that the scheduler is being initialized with. |
|
|
|
Returns: |
|
`torch.FloatTensor`: rescaled betas with zero terminal SNR |
|
""" |
|
|
|
alphas = 1.0 - betas |
|
alphas_cumprod = torch.cumprod(alphas, dim=0) |
|
alphas_bar_sqrt = alphas_cumprod.sqrt() |
|
|
|
|
|
alphas_bar_sqrt_0 = alphas_bar_sqrt[0].clone() |
|
alphas_bar_sqrt_T = alphas_bar_sqrt[-1].clone() |
|
|
|
|
|
alphas_bar_sqrt -= alphas_bar_sqrt_T |
|
|
|
|
|
alphas_bar_sqrt *= alphas_bar_sqrt_0 / (alphas_bar_sqrt_0 - alphas_bar_sqrt_T) |
|
|
|
|
|
alphas_bar = alphas_bar_sqrt**2 |
|
alphas = alphas_bar[1:] / alphas_bar[:-1] |
|
alphas = torch.cat([alphas_bar[0:1], alphas]) |
|
betas = 1 - alphas |
|
|
|
return betas |
|
|
|
|
|
class LCMSingleStepScheduler(SchedulerMixin, ConfigMixin): |
|
""" |
|
`LCMSingleStepScheduler` extends the denoising procedure introduced in denoising diffusion probabilistic models (DDPMs) with |
|
non-Markovian guidance. |
|
|
|
This model inherits from [`SchedulerMixin`] and [`ConfigMixin`]. [`~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`, 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`. |
|
trained_betas (`np.ndarray`, *optional*): |
|
Pass an array of betas directly to the constructor to bypass `beta_start` and `beta_end`. |
|
original_inference_steps (`int`, *optional*, defaults to 50): |
|
The default number of inference steps used to generate a linearly-spaced timestep schedule, from which we |
|
will ultimately take `num_inference_steps` evenly spaced timesteps to form the final timestep schedule. |
|
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`. |
|
set_alpha_to_one (`bool`, defaults to `True`): |
|
Each diffusion step uses the alphas product value at that step and at the previous one. For the final step |
|
there is no previous alpha. When this option is `True` the previous alpha product is fixed to `1`, |
|
otherwise it uses the alpha value at step 0. |
|
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. |
|
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. |
|
timestep_scaling (`float`, defaults to 10.0): |
|
The factor the timesteps will be multiplied by when calculating the consistency model boundary conditions |
|
`c_skip` and `c_out`. Increasing this will decrease the approximation error (although the approximation |
|
error at the default of `10.0` is already pretty small). |
|
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). |
|
""" |
|
|
|
order = 1 |
|
|
|
@register_to_config |
|
def __init__( |
|
self, |
|
num_train_timesteps: int = 1000, |
|
beta_start: float = 0.00085, |
|
beta_end: float = 0.012, |
|
beta_schedule: str = "scaled_linear", |
|
trained_betas: Optional[Union[np.ndarray, List[float]]] = None, |
|
original_inference_steps: int = 50, |
|
clip_sample: bool = False, |
|
clip_sample_range: float = 1.0, |
|
set_alpha_to_one: bool = True, |
|
steps_offset: int = 0, |
|
prediction_type: str = "epsilon", |
|
thresholding: bool = False, |
|
dynamic_thresholding_ratio: float = 0.995, |
|
sample_max_value: float = 1.0, |
|
timestep_spacing: str = "leading", |
|
timestep_scaling: float = 10.0, |
|
rescale_betas_zero_snr: bool = False, |
|
): |
|
if trained_betas is not None: |
|
self.betas = torch.tensor(trained_betas, dtype=torch.float32) |
|
elif beta_schedule == "linear": |
|
self.betas = torch.linspace(beta_start, beta_end, num_train_timesteps, dtype=torch.float32) |
|
elif beta_schedule == "scaled_linear": |
|
|
|
self.betas = ( |
|
torch.linspace(beta_start**0.5, beta_end**0.5, num_train_timesteps, dtype=torch.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__}") |
|
|
|
|
|
if rescale_betas_zero_snr: |
|
self.betas = rescale_zero_terminal_snr(self.betas) |
|
|
|
self.alphas = 1.0 - self.betas |
|
self.alphas_cumprod = torch.cumprod(self.alphas, dim=0) |
|
|
|
|
|
|
|
|
|
|
|
self.final_alpha_cumprod = torch.tensor(1.0) if set_alpha_to_one else self.alphas_cumprod[0] |
|
|
|
|
|
self.init_noise_sigma = 1.0 |
|
|
|
|
|
self.num_inference_steps = None |
|
self.timesteps = torch.from_numpy(np.arange(0, num_train_timesteps)[::-1].copy().astype(np.int64)) |
|
|
|
self._step_index = None |
|
|
|
|
|
def _init_step_index(self, timestep): |
|
if isinstance(timestep, torch.Tensor): |
|
timestep = timestep.to(self.timesteps.device) |
|
|
|
index_candidates = (self.timesteps == timestep).nonzero() |
|
|
|
|
|
|
|
|
|
|
|
if len(index_candidates) > 1: |
|
step_index = index_candidates[1] |
|
else: |
|
step_index = index_candidates[0] |
|
|
|
self._step_index = step_index.item() |
|
|
|
@property |
|
def step_index(self): |
|
return self._step_index |
|
|
|
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: |
|
sample (`torch.FloatTensor`): |
|
The input sample. |
|
timestep (`int`, *optional*): |
|
The current timestep in the diffusion chain. |
|
Returns: |
|
`torch.FloatTensor`: |
|
A scaled input sample. |
|
""" |
|
return sample |
|
|
|
|
|
def _threshold_sample(self, sample: torch.FloatTensor) -> torch.FloatTensor: |
|
""" |
|
"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 |
|
batch_size, channels, *remaining_dims = sample.shape |
|
|
|
if dtype not in (torch.float32, torch.float64): |
|
sample = sample.float() |
|
|
|
|
|
sample = sample.reshape(batch_size, channels * np.prod(remaining_dims)) |
|
|
|
abs_sample = sample.abs() |
|
|
|
s = torch.quantile(abs_sample, self.config.dynamic_thresholding_ratio, dim=1) |
|
s = torch.clamp( |
|
s, min=1, max=self.config.sample_max_value |
|
) |
|
s = s.unsqueeze(1) |
|
sample = torch.clamp(sample, -s, s) / s |
|
|
|
sample = sample.reshape(batch_size, channels, *remaining_dims) |
|
sample = sample.to(dtype) |
|
|
|
return sample |
|
|
|
def set_timesteps( |
|
self, |
|
num_inference_steps: int = None, |
|
device: Union[str, torch.device] = None, |
|
original_inference_steps: Optional[int] = None, |
|
strength: int = 1.0, |
|
timesteps: Optional[list] = None, |
|
): |
|
""" |
|
Sets the discrete timesteps used for the diffusion chain (to be run before inference). |
|
|
|
Args: |
|
num_inference_steps (`int`): |
|
The number of diffusion steps used when generating samples with a pre-trained model. |
|
device (`str` or `torch.device`, *optional*): |
|
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. |
|
original_inference_steps (`int`, *optional*): |
|
The original number of inference steps, which will be used to generate a linearly-spaced timestep |
|
schedule (which is different from the standard `diffusers` implementation). We will then take |
|
`num_inference_steps` timesteps from this schedule, evenly spaced in terms of indices, and use that as |
|
our final timestep schedule. If not set, this will default to the `original_inference_steps` attribute. |
|
""" |
|
|
|
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) |
|
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." |
|
) |
|
|
|
self.num_inference_steps = num_inference_steps |
|
original_steps = ( |
|
original_inference_steps if original_inference_steps is not None else self.config.original_inference_steps |
|
) |
|
|
|
if original_steps > self.config.num_train_timesteps: |
|
raise ValueError( |
|
f"`original_steps`: {original_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." |
|
) |
|
|
|
if num_inference_steps > original_steps: |
|
raise ValueError( |
|
f"`num_inference_steps`: {num_inference_steps} cannot be larger than `original_inference_steps`:" |
|
f" {original_steps} because the final timestep schedule will be a subset of the" |
|
f" `original_inference_steps`-sized initial timestep schedule." |
|
) |
|
|
|
|
|
|
|
c = self.config.num_train_timesteps // original_steps |
|
|
|
lcm_origin_timesteps = np.asarray(list(range(1, int(original_steps * strength) + 1))) * c - 1 |
|
skipping_step = len(lcm_origin_timesteps) // num_inference_steps |
|
|
|
timesteps = lcm_origin_timesteps[::-skipping_step][:num_inference_steps] |
|
|
|
self.timesteps = torch.from_numpy(timesteps.copy()).to(device=device, dtype=torch.long) |
|
|
|
self._step_index = None |
|
|
|
def get_scalings_for_boundary_condition_discrete(self, timestep): |
|
self.sigma_data = 0.5 |
|
scaled_timestep = timestep * self.config.timestep_scaling |
|
|
|
c_skip = self.sigma_data**2 / (scaled_timestep**2 + self.sigma_data**2) |
|
c_out = scaled_timestep / (scaled_timestep**2 + self.sigma_data**2) ** 0.5 |
|
return c_skip, c_out |
|
|
|
def append_dims(self, x, target_dims): |
|
"""Appends dimensions to the end of a tensor until it has target_dims dimensions.""" |
|
dims_to_append = target_dims - x.ndim |
|
if dims_to_append < 0: |
|
raise ValueError(f"input has {x.ndim} dims but target_dims is {target_dims}, which is less") |
|
return x[(...,) + (None,) * dims_to_append] |
|
|
|
def extract_into_tensor(self, a, t, x_shape): |
|
b, *_ = t.shape |
|
out = a.gather(-1, t) |
|
return out.reshape(b, *((1,) * (len(x_shape) - 1))) |
|
|
|
def step( |
|
self, |
|
model_output: torch.FloatTensor, |
|
timestep: torch.Tensor, |
|
sample: torch.FloatTensor, |
|
generator: Optional[torch.Generator] = None, |
|
return_dict: bool = True, |
|
) -> Union[LCMSingleStepSchedulerOutput, Tuple]: |
|
""" |
|
Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion |
|
process from the learned model outputs (most often the predicted noise). |
|
|
|
Args: |
|
model_output (`torch.FloatTensor`): |
|
The direct output from learned diffusion model. |
|
timestep (`float`): |
|
The current discrete timestep in the diffusion chain. |
|
sample (`torch.FloatTensor`): |
|
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_lcm.LCMSchedulerOutput`] or `tuple`. |
|
Returns: |
|
[`~schedulers.scheduling_utils.LCMSchedulerOutput`] or `tuple`: |
|
If return_dict is `True`, [`~schedulers.scheduling_lcm.LCMSchedulerOutput`] is returned, otherwise a |
|
tuple is returned where the first element is the sample tensor. |
|
""" |
|
|
|
alphas_cumprod = self.alphas_cumprod.to(sample.device) |
|
|
|
|
|
if timestep.ndim == 0: |
|
timestep = timestep.unsqueeze(0) |
|
alpha_prod_t = self.extract_into_tensor(alphas_cumprod, timestep, sample.shape) |
|
beta_prod_t = 1 - alpha_prod_t |
|
|
|
|
|
c_skip, c_out = self.get_scalings_for_boundary_condition_discrete(timestep) |
|
c_skip, c_out = [self.append_dims(x, sample.ndim) for x in [c_skip, c_out]] |
|
|
|
|
|
if self.config.prediction_type == "epsilon": |
|
predicted_original_sample = (sample - torch.sqrt(beta_prod_t) * model_output) / torch.sqrt(alpha_prod_t) |
|
elif self.config.prediction_type == "sample": |
|
predicted_original_sample = model_output |
|
elif self.config.prediction_type == "v_prediction": |
|
predicted_original_sample = torch.sqrt(alpha_prod_t) * sample - torch.sqrt(beta_prod_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 `LCMScheduler`." |
|
) |
|
|
|
|
|
if self.config.thresholding: |
|
predicted_original_sample = self._threshold_sample(predicted_original_sample) |
|
elif self.config.clip_sample: |
|
predicted_original_sample = predicted_original_sample.clamp( |
|
-self.config.clip_sample_range, self.config.clip_sample_range |
|
) |
|
|
|
|
|
denoised = c_out * predicted_original_sample + c_skip * sample |
|
|
|
if not return_dict: |
|
return (denoised, ) |
|
|
|
return LCMSingleStepSchedulerOutput(denoised=denoised) |
|
|
|
|
|
def add_noise( |
|
self, |
|
original_samples: torch.FloatTensor, |
|
noise: torch.FloatTensor, |
|
timesteps: torch.IntTensor, |
|
) -> torch.FloatTensor: |
|
|
|
alphas_cumprod = self.alphas_cumprod.to(device=original_samples.device, dtype=original_samples.dtype) |
|
timesteps = timesteps.to(original_samples.device) |
|
|
|
sqrt_alpha_prod = 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 - 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 get_velocity( |
|
self, sample: torch.FloatTensor, noise: torch.FloatTensor, timesteps: torch.IntTensor |
|
) -> torch.FloatTensor: |
|
|
|
alphas_cumprod = self.alphas_cumprod.to(device=sample.device, dtype=sample.dtype) |
|
timesteps = timesteps.to(sample.device) |
|
|
|
sqrt_alpha_prod = alphas_cumprod[timesteps] ** 0.5 |
|
sqrt_alpha_prod = sqrt_alpha_prod.flatten() |
|
while len(sqrt_alpha_prod.shape) < len(sample.shape): |
|
sqrt_alpha_prod = sqrt_alpha_prod.unsqueeze(-1) |
|
|
|
sqrt_one_minus_alpha_prod = (1 - 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(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 |
|
|
|
def __len__(self): |
|
return self.config.num_train_timesteps |
|
|