diffusion_models.gaussian_diffusion.beta_schedulers

Module Contents

class BaseBetaScheduler(steps, enforce_zero_terminal_snr=False, initialize=True)[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.
initialize (bool) – Whether to initialize the beta scheduler. If this is set to False, you will need to manually set self.betas and self.alpha_bars. Otherwise, they are initialized using your sample_betas and compute_alpha_bar methods.

Warning

Do not instantiate this class directly. Instead, build your own Beta scheduler by inheriting from BaseBetaScheduler. (see LinearBetaScheduler)

steps: int[source]: The number of steps for the beta scheduler.

betas = None[source]: The \(\beta\) computed according to sample_betas().

alpha_bars = None[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:: Tensor

abstract compute_alpha_bar()[source]

Compute \(\bar{\alpha}\) for noise scheduling.

Returns:: A torch tensor of the \(\bar{\alpha}\) values.
Return type:: Tensor

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”.

classmethod from_tensors(steps, betas, alpha_bars)[source]

Instantiate a beta scheduler from tensors.

Instantiate a beta scheduler from tensors. This is particularly useful for loading checkpoints.

Parameters:

steps (int) – The number of steps for the beta scheduler.
betas (Tensor) – The pre-computed beta values for the noise scheduler.
alpha_bars (Tensor) – The pre-computed alpha bar values for the noise scheduler.

Returns:

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_start and beta_end.

Parameters:

beta_start (float) – The starting value of the betas.
beta_end (float) – The end value of the betas.
steps (int) – The number of steps for the beta scheduler. This is also the number of betas.
enforce_zero_terminal_snr (bool) – Whether to enforce zero terminal SNR.

beta_start: int[source]: The starting value of the betas.

beta_end: int[source]: The end value of the betas.

sample_betas()[source]

Return linearly spaced betas between self.beta_start and self.beta_end.

compute_alpha_bar()[source]: Return \(\bar{\alpha}\) computed from the beta values.

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:

offset (float) – The offset \(s\) defined above.
steps (int) – The number of steps for the beta scheduler.
max_beta (Optional[float]) – The maximum beta values. Higher values will be clipped.

offset: float[source]: The offset \(s\) defined above.

max_beta: float | None[source]: The maximum beta values. Higher values will be clipped.

steps: int[source]: The number of steps for the beta scheduler.

f(t)[source]

A helper function to compute the \(\bar{\alpha}_t\).

Parameters:

t (Tensor) – The timestep to compute.

Returns:

\[f(t) = \cos(\frac{t/T + s}{1 + s} * \frac{\pi}{2})^2\]

Return type:

Tensor

sample_betas()[source]

Compute \(\beta\) for noise scheduling.

Returns:: A torch tensor of the \(\beta\) values.

compute_alpha_bar()[source]

Compute \(\bar{\alpha}\) for noise scheduling.

Returns:: A torch tensor of the \(\bar{\alpha}\) values.