|
<!--Copyright 2023 The HuggingFace Team. All rights reserved. |
|
|
|
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. |
|
--> |
|
|
|
# Schedulers |
|
|
|
Diffusers contains multiple pre-built schedule functions for the diffusion process. |
|
|
|
## What is a scheduler? |
|
|
|
The schedule functions, denoted *Schedulers* in the library take in the output of a trained model, a sample which the diffusion process is iterating on, and a timestep to return a denoised sample. That's why schedulers may also be called *Samplers* in other diffusion models implementations. |
|
|
|
- Schedulers define the methodology for iteratively adding noise to an image or for updating a sample based on model outputs. |
|
- adding noise in different manners represent the algorithmic processes to train a diffusion model by adding noise to images. |
|
- for inference, the scheduler defines how to update a sample based on an output from a pretrained model. |
|
- Schedulers are often defined by a *noise schedule* and an *update rule* to solve the differential equation solution. |
|
|
|
### Discrete versus continuous schedulers |
|
|
|
All schedulers take in a timestep to predict the updated version of the sample being diffused. |
|
The timesteps dictate where in the diffusion process the step is, where data is generated by iterating forward in time and inference is executed by propagating backwards through timesteps. |
|
Different algorithms use timesteps that can be discrete (accepting `int` inputs), such as the [`DDPMScheduler`] or [`PNDMScheduler`], or continuous (accepting `float` inputs), such as the score-based schedulers [`ScoreSdeVeScheduler`] or [`ScoreSdeVpScheduler`]. |
|
|
|
## Designing Re-usable schedulers |
|
|
|
The core design principle between the schedule functions is to be model, system, and framework independent. |
|
This allows for rapid experimentation and cleaner abstractions in the code, where the model prediction is separated from the sample update. |
|
To this end, the design of schedulers is such that: |
|
|
|
- Schedulers can be used interchangeably between diffusion models in inference to find the preferred trade-off between speed and generation quality. |
|
- Schedulers are currently by default in PyTorch, but are designed to be framework independent (partial Jax support currently exists). |
|
- Many diffusion pipelines, such as [`StableDiffusionPipeline`] and [`DiTPipeline`] can use any of [`KarrasDiffusionSchedulers`] |
|
|
|
## Schedulers Summary |
|
|
|
The following table summarizes all officially supported schedulers, their corresponding paper |
|
|
|
| Scheduler | Paper | |
|
|---|---| |
|
| [ddim](./ddim) | [**Denoising Diffusion Implicit Models**](https://arxiv.org/abs/2010.02502) | |
|
| [ddim_inverse](./ddim_inverse) | [**Denoising Diffusion Implicit Models**](https://arxiv.org/abs/2010.02502) | |
|
| [ddpm](./ddpm) | [**Denoising Diffusion Probabilistic Models**](https://arxiv.org/abs/2006.11239) | |
|
| [deis](./deis) | [**DEISMultistepScheduler**](https://arxiv.org/abs/2204.13902) | |
|
| [singlestep_dpm_solver](./singlestep_dpm_solver) | [**Singlestep DPM-Solver**](https://arxiv.org/abs/2206.00927) | |
|
| [multistep_dpm_solver](./multistep_dpm_solver) | [**Multistep DPM-Solver**](https://arxiv.org/abs/2206.00927) | |
|
| [heun](./heun) | [**Heun scheduler inspired by Karras et. al paper**](https://arxiv.org/abs/2206.00364) | |
|
| [dpm_discrete](./dpm_discrete) | [**DPM Discrete Scheduler inspired by Karras et. al paper**](https://arxiv.org/abs/2206.00364) | |
|
| [dpm_discrete_ancestral](./dpm_discrete_ancestral) | [**DPM Discrete Scheduler with ancestral sampling inspired by Karras et. al paper**](https://arxiv.org/abs/2206.00364) | |
|
| [stochastic_karras_ve](./stochastic_karras_ve) | [**Variance exploding, stochastic sampling from Karras et. al**](https://arxiv.org/abs/2206.00364) | |
|
| [lms_discrete](./lms_discrete) | [**Linear multistep scheduler for discrete beta schedules**](https://arxiv.org/abs/2206.00364) | |
|
| [pndm](./pndm) | [**Pseudo numerical methods for diffusion models (PNDM)**](https://github.com/crowsonkb/k-diffusion/blob/481677d114f6ea445aa009cf5bd7a9cdee909e47/k_diffusion/sampling.py#L181) | |
|
| [score_sde_ve](./score_sde_ve) | [**variance exploding stochastic differential equation (VE-SDE) scheduler**](https://arxiv.org/abs/2011.13456) | |
|
| [ipndm](./ipndm) | [**improved pseudo numerical methods for diffusion models (iPNDM)**](https://github.com/crowsonkb/v-diffusion-pytorch/blob/987f8985e38208345c1959b0ea767a625831cc9b/diffusion/sampling.py#L296) | |
|
| [score_sde_vp](./score_sde_vp) | [**Variance preserving stochastic differential equation (VP-SDE) scheduler**](https://arxiv.org/abs/2011.13456) | |
|
| [euler](./euler) | [**Euler scheduler**](https://arxiv.org/abs/2206.00364) | |
|
| [euler_ancestral](./euler_ancestral) | [**Euler Ancestral scheduler**](https://github.com/crowsonkb/k-diffusion/blob/481677d114f6ea445aa009cf5bd7a9cdee909e47/k_diffusion/sampling.py#L72) | |
|
| [vq_diffusion](./vq_diffusion) | [**VQDiffusionScheduler**](https://arxiv.org/abs/2111.14822) | |
|
| [unipc](./unipc) | [**UniPCMultistepScheduler**](https://arxiv.org/abs/2302.04867) | |
|
| [repaint](./repaint) | [**RePaint scheduler**](https://arxiv.org/abs/2201.09865) | |
|
|
|
## API |
|
|
|
The core API for any new scheduler must follow a limited structure. |
|
- Schedulers should provide one or more `def step(...)` functions that should be called to update the generated sample iteratively. |
|
- Schedulers should provide a `set_timesteps(...)` method that configures the parameters of a schedule function for a specific inference task. |
|
- Schedulers should be framework-specific. |
|
|
|
The base class [`SchedulerMixin`] implements low level utilities used by multiple schedulers. |
|
|
|
### SchedulerMixin |
|
[[autodoc]] SchedulerMixin |
|
|
|
### SchedulerOutput |
|
The class [`SchedulerOutput`] contains the outputs from any schedulers `step(...)` call. |
|
|
|
[[autodoc]] schedulers.scheduling_utils.SchedulerOutput |
|
|
|
### KarrasDiffusionSchedulers |
|
|
|
`KarrasDiffusionSchedulers` encompasses the main generalization of schedulers in Diffusers. The schedulers in this class are distinguished, at a high level, by their noise sampling strategy; the type of network and scaling; and finally the training strategy or how the loss is weighed. |
|
|
|
The different schedulers, depending on the type of ODE solver, fall into the above taxonomy and provide a good abstraction for the design of the main schedulers implemented in Diffusers. The schedulers in this class are given below: |
|
|
|
[[autodoc]] schedulers.scheduling_utils.KarrasDiffusionSchedulers |
|
|