diffusion_models.gaussian_diffusion.beta_schedulers
Module Contents
- class BaseBetaScheduler(steps, enforce_zero_terminal_snr=False)[source]
Initializes a beta scheduler.
BaseBetaScheduler is an abstract base class for different beta scheduler implementations. It defines the interface that all beta schedulers should adhere to.
- Parameters:
steps (int) – The number of steps for the beta.
enforce_zero_terminal_snr (bool) –
Whether to enforce zero terminal SNR inline with “Common Diffusion Noise Schedules and Sample Steps are Flawed”.
Defaults to
False.
Warning
Do not instantiate this class directly. Instead, build your own Beta scheduler by inheriting from BaseBetaScheduler. (see
LinearBetaScheduler)- betas[source]
The \(\beta\) computed according to
sample_betas().
- alpha_bars[source]
The \(\bar{\alpha}\) computed according to
compute_alpha_bar().
- enforce_zero_terminal_snr()[source]
Enforce terminal SNR by adjusting \(\beta\) and \(\bar{\alpha}\).
This method enforces zero terminal SNR according to “Common Diffusion Noise Schedules and Sample Steps are Flawed”.
- abstract sample_betas()[source]
Compute \(\beta\) for noise scheduling.
- Returns:
A torch tensor of the \(\beta\) values.
- Return type:
- abstract compute_alpha_bar()[source]
Compute \(\bar{\alpha}\) for noise scheduling.
- Returns:
A torch tensor of the \(\bar{\alpha}\) values.
- Return type:
- to(device)[source]
Moves the beta scheduler to the given device.
- Parameters:
device (str) – The device to which the method should move the object. Default is “cpu”.
- class LinearBetaScheduler(beta_start=0.0001, beta_end=0.02, steps=1000, enforce_zero_terminal_snr=True)[source]
A Linear Beta scheduler.
A simple linear beta scheduler with betas linearly spaced between
beta_startandbeta_end.- Parameters:
- class CosineBetaScheduler(offset=0.008, steps=1000, max_beta=0.999)[source]
A Cosine Beta scheduler.
A Cosine Beta Scheduler based on the following formulas:
\begin{equation} \left\{ \begin{aligned} \bar{\alpha}_t &= \frac{f(t)}{f(0)} \\ \beta_t &= 1 - \frac{\bar{\alpha}_t}{\bar{\alpha}_t -1} \end{aligned} \right. \end{equation}where
\[f(t) = \cos(\frac{t/T + s}{1 + s} * \frac{\pi}{2})^2\]where
\begin{equation} \left\{ \begin{aligned} s & \text{ is the offset} \\ T & \text{ is the number of steps} \end{aligned} \right. \end{equation}- Parameters: