DiffusionScheduler

class DiffusionScheduler(alpha_bar: Tensor)[source]

Bases: Scheduler[DDTensor]

Scheduler for discrete-time diffusion models based on a given noise schedule.

Parameters:

alpha_bar (torch.Tensor) – Cumulative product of (1 - beta) values, shape (K,), where K is the number of diffusion steps.

alpha(x: DDTensor, t: Tensor) → DDTensor[source]

\(\alpha_t\).

Can be overwritten if \(\alpha_t\) is data-dependent.

Parameters:

  • x (D) – Data tensor.

  • t (torch.Tensor, shape (n,)) – Time tensor with values in [0, 1].

Returns:

alpha – Values of \(\alpha_t\) at the given times.

Return type:

D, same data shape as x

alpha_dot(x: DDTensor, t: Tensor) → DDTensor[source]

\(\dot{\alpha}_t\).

Can be overwritten if \(\dot{\alpha}_t\) is data-dependent.

Parameters:

  • x (D) – Data tensor.

  • t (torch.Tensor, shape (n,)) – Time tensor with values in [0, 1].

Returns:

alpha_dot – Values of \(\dot{\alpha}_t\) at the given times.

Return type:

D, same data shape as x

beta(x: DDTensor, t: Tensor) → DDTensor[source]

\(\beta_t\).

Parameters:

  • x (D) – Data tensor.

  • t (torch.Tensor, shape (n,)) – Time tensor with values in [0, 1].

Returns:

beta – Values of \(\beta_t\) at the given times.

Return type:

D, same data shape as x

beta_dot(x: DDTensor, t: Tensor) → DDTensor[source]

\(\dot{\beta}_t\).

Can be overwritten if \(\dot{\beta}_t\) is data-dependent.

Parameters:

  • x (D) – Data tensor.

  • t (torch.Tensor, shape (n,)) – Time tensor with values in [0, 1].

Returns:

beta_dot – Values of \(\dot{\beta}_t\) at the given times.

Return type:

D, same data shape as x

model_input(t: Tensor) → Tensor[source]

Input to the model at time t that encodes the timestep.