schedulers.mdx 7.6 KB
Newer Older
Nathan Lambert's avatar
Nathan Lambert committed
1
2
3
4
5
6
7
8
9
10
11
12
<!--Copyright 2022 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.
-->

13
14
15
16
# Schedulers

Diffusers contains multiple pre-built schedule functions for the diffusion process.

17
18
## What is a scheduler?

19
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.
20
21
22
23
24
25
26

- 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
27

28
29
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.
30
Different algorithms use timesteps that both discrete (accepting `int` inputs), such as the [`DDPMScheduler`] or [`PNDMScheduler`], and continuous (accepting `float` inputs), such as the score-based schedulers [`ScoreSdeVeScheduler`] or [`ScoreSdeVpScheduler`].
31
32

## Designing Re-usable schedulers
33

34
35
36
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:
37

38
- Schedulers can be used interchangeably between diffusion models in inference to find the preferred trade-off between speed and generation quality.
39
- Schedulers are currently by default in PyTorch, but are designed to be framework independent (partial Jax support currently exists).
Nathan Lambert's avatar
Nathan Lambert committed
40
41
42


## API
43

44
45
46
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.
47
- Schedulers should be framework-specific.
48
49
50

The base class [`SchedulerMixin`] implements low level utilities used by multiple schedulers.

51
### SchedulerMixin
52
53
[[autodoc]] SchedulerMixin

54
### SchedulerOutput
55
The class [`SchedulerOutput`] contains the outputs from any schedulers `step(...)` call.
56

57
58
[[autodoc]] schedulers.scheduling_utils.SchedulerOutput

59
### Implemented Schedulers
60
61
62
63
64

#### Denoising diffusion implicit models (DDIM)

Original paper can be found here.

65
[[autodoc]] DDIMScheduler
66
67
68
69
70

#### Denoising diffusion probabilistic models (DDPM)

Original paper can be found [here](https://arxiv.org/abs/2010.02502).

71
[[autodoc]] DDPMScheduler
72

73
74
75
76
77
78
#### Multistep DPM-Solver

Original paper can be found [here](https://arxiv.org/abs/2206.00927) and the [improved version](https://arxiv.org/abs/2211.01095). The original implementation can be found [here](https://github.com/LuChengTHU/dpm-solver).

[[autodoc]] DPMSolverMultistepScheduler

79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
#### Heun scheduler inspired by Karras et. al paper

Algorithm 1 of [Karras et. al](https://arxiv.org/abs/2206.00364).
Scheduler ported from @crowsonkb's https://github.com/crowsonkb/k-diffusion library:

All credit for making this scheduler work goes to [Katherine Crowson](https://github.com/crowsonkb/)

[[autodoc]] HeunDiscreteScheduler

#### DPM Discrete Scheduler inspired by Karras et. al paper

Inspired by [Karras et. al](https://arxiv.org/abs/2206.00364).
Scheduler ported from @crowsonkb's https://github.com/crowsonkb/k-diffusion library:

All credit for making this scheduler work goes to [Katherine Crowson](https://github.com/crowsonkb/)

[[autodoc]] KDPM2DiscreteScheduler

#### DPM Discrete Scheduler with ancestral sampling inspired by Karras et. al paper

Inspired by [Karras et. al](https://arxiv.org/abs/2206.00364).
Scheduler ported from @crowsonkb's https://github.com/crowsonkb/k-diffusion library:

All credit for making this scheduler work goes to [Katherine Crowson](https://github.com/crowsonkb/)

[[autodoc]] KDPM2AncestralDiscreteScheduler

106
#### Variance exploding, stochastic sampling from Karras et. al
107
108
109

Original paper can be found [here](https://arxiv.org/abs/2006.11239).

110
[[autodoc]] KarrasVeScheduler
111
112
113
114
115

#### Linear multistep scheduler for discrete beta schedules

Original implementation can be found [here](https://arxiv.org/abs/2206.00364).

116
[[autodoc]] LMSDiscreteScheduler
117
118
119
120
121

#### Pseudo numerical methods for diffusion models (PNDM)

Original implementation can be found [here](https://github.com/crowsonkb/k-diffusion/blob/481677d114f6ea445aa009cf5bd7a9cdee909e47/k_diffusion/sampling.py#L181).

122
[[autodoc]] PNDMScheduler
123

124
#### variance exploding stochastic differential equation (VE-SDE) scheduler
125
126
127

Original paper can be found [here](https://arxiv.org/abs/2011.13456).

128
[[autodoc]] ScoreSdeVeScheduler
129

130
131
132
133
#### improved pseudo numerical methods for diffusion models (iPNDM)

Original implementation can be found [here](https://github.com/crowsonkb/v-diffusion-pytorch/blob/987f8985e38208345c1959b0ea767a625831cc9b/diffusion/sampling.py#L296).

134
135
136
[[autodoc]] IPNDMScheduler

#### variance preserving stochastic differential equation (VP-SDE) scheduler
137
138
139
140

Original paper can be found [here](https://arxiv.org/abs/2011.13456).

<Tip warning={true}>
Nathan Lambert's avatar
Nathan Lambert committed
141

142
Score SDE-VP is under construction.
Nathan Lambert's avatar
Nathan Lambert committed
143

144
</Tip>
Nathan Lambert's avatar
Nathan Lambert committed
145

146
[[autodoc]] schedulers.scheduling_sde_vp.ScoreSdeVpScheduler
147
148
149
150
151
152
153
154
155
156
157
158
159
160

#### Euler scheduler

Euler scheduler (Algorithm 2) from the paper [Elucidating the Design Space of Diffusion-Based Generative Models](https://arxiv.org/abs/2206.00364) by Karras et al. (2022). Based on the original [k-diffusion](https://github.com/crowsonkb/k-diffusion/blob/481677d114f6ea445aa009cf5bd7a9cdee909e47/k_diffusion/sampling.py#L51) implementation by Katherine Crowson.
Fast scheduler which often times generates good outputs with 20-30 steps.

[[autodoc]] EulerDiscreteScheduler


#### Euler Ancestral scheduler

Ancestral sampling with Euler method steps. Based on the original (k-diffusion)[https://github.com/crowsonkb/k-diffusion/blob/481677d114f6ea445aa009cf5bd7a9cdee909e47/k_diffusion/sampling.py#L72] implementation by Katherine Crowson.
Fast scheduler which often times generates good outputs with 20-30 steps.

Revist's avatar
Revist committed
161
162
163
[[autodoc]] EulerAncestralDiscreteScheduler


Will Berman's avatar
Will Berman committed
164
165
166
167
168
169
#### VQDiffusionScheduler

Original paper can be found [here](https://arxiv.org/abs/2111.14822)

[[autodoc]] VQDiffusionScheduler

Revist's avatar
Revist committed
170
171
172
173
174
175
176
#### RePaint scheduler

DDPM-based inpainting scheduler for unsupervised inpainting with extreme masks. 
Intended for use with [`RePaintPipeline`].
Based on the paper [RePaint: Inpainting using Denoising Diffusion Probabilistic Models](https://arxiv.org/abs/2201.09865) 
and the original implementation by Andreas Lugmayr et al.: https://github.com/andreas128/RePaint

Will Berman's avatar
Will Berman committed
177
[[autodoc]] RePaintScheduler