Extract values from a 1-D numpy array for a batch of indices. :param arr: the 1-D numpy array. :param timesteps: a tensor of indices into

Create a 1D, 2D, or 3D convolution module.

Function to reshape the vector by the number of dimensions of x. E.g. x (bs, c, h, w), v (bs) -> v (bs, 1, 1, 1).

Take the mean over all non-batch dimensions.

Diffuse the data for a given number of diffusion steps => sample x_t ~ q(x_t | x_0) with x_t = sqrt(alpha_bar_t) *

Create a linear module.

Make a standard normalization layer. :param channels: number of input channels. :return: an nn.Module for normalization.

Compute the noise coefficient along the path

Take the mean over all non-batch dimensions.

Apply the model to get p(x_{t-1} | x_t), as well as a prediction of the initial x, x_0. :param model: the model, which takes

Compute the mean and variance of the diffusion posterior: q(x_{t-1} | x_t, x_0)

Get a term for the variational lower-bound. The resulting units are bits (rather than nats, as one might expect). This allows

Evaluate a function without caching intermediate activations, allowing for reduced memory at the expense of extra compute in the backward pas

Compute the time-dependent conditional vector field ut = alpha_dt_t * x1 + sigma_dt_t * x0, see Eq. (7) in [3].

Sample from the time-dependent density p_t xt ~ N(alpha_t * x1 + sigma_t * x0, sigma_min * I), according to Eq. (1) in [3

Takes batch and extracts data. Returns: x0: Samples from source distribution (can also be None if we start from noise).

Sample the v-parameterized vector field at time t

Create sinusoidal timestep embeddings. :param timesteps: a 1-D Tensor of N indices, one per batch element. These may be

Zero out the parameters of a module and return it.

kwargs must contain the output keys with their corresponding shapes, e.g. image=(3, 32, 32), label=(1,)

A fast approximation of the cumulative distribution function of the standard normal.

A fast approximation of the cumulative distribution function of the standard normal.

We always output sde according to score parametrization;

Compute what the p_mean_variance output would have been, should the model's score function be conditioned by cond_fn. See con

Convert the continuous time t in [0,1] to discrete time t [0, 1000] # TODO: Make it compatible with zero-terminal SNR

A counter for the `thop` package to count the operations in an attention operation. Meant to be used like: macs, params = thop.pr

Function to include sampling with Classifier-Free Guidance (CFG)

Compute the KL divergence between two gaussians.

Compute the KL divergence between two gaussians. Shapes are automatically broadcasted, so batches can be compared to scalars, among other

Apply the model to get p(x_{t-1} | x_t), as well as a prediction of the initial x_0.

Function to reshape the vector by the number of dimensions of x. E.g. x (bs, c, h, w), v (bs) -> v (bs, 1, 1, 1). Args: x :

Normalize each sample in a batch independently with min-max normalization to [0, 1]

Compute the mean and variance of the diffusion posterior q(x_{t-1} | x_t, x_0)

Diffuse the data for a given number of diffusion steps. In other words, sample from q(x_t | x_0). :param x_start: the initial

Args: - sampling_method: type of sampler used in solving the SDE; default to be Euler-Maruyama - diffusion_form: function for

Args: tensor: Tensor of shape (..., 3, h, w) in range [-1, 1] or [0, 1] denormalize_zero_one: If True, denormalizes image from ra

Convert from [-1, 1] to [0, 255]

Zero out the parameters of a module and return it.

Get the prior KL term for the variational lower-bound, measured in bits-per-dim. This term can't be optimized, as it only dep