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 int

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

Defaults for image training.

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

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

Write the sequence of args, with no separators, to the console and output files (if you've configured an output file).

Take the mean over all non-batch dimensions.

Create a linear module.

Convert the torso of the model to float16.

For a dataset, create a generator over (images, kwargs) pairs. Each images is an NCHW float tensor, and the kwargs dict contains zero or

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)

Zero out the parameters of a module and return it.

Defaults for classifier models.

Use DDIM to sample from the model and yield intermediate samples from each timestep of DDIM. Same usage as p_sample_loop_pro

Generate samples from the model and yield intermediate samples from each timestep of diffusion. Arguments are the same as p_

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

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

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

If comm is provided, average all numerical stats across that comm

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

Create a ScheduleSampler from a library of pre-defined samplers. :param name: the name of the sampler. :param diffusion: the diffusion o

Defaults for image and classifier training.

Copy model parameters into a (differently-shaped) list of full-precision parameters.

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

Parse filenames of the form path/to/modelNNNNNN.pt, where NNNNNN is the checkpoint's number of steps.

Importance-sample timesteps for a batch. :param batch_size: the number of timesteps. :param device: the torch device to save

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

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

Create a 1D, 2D, or 3D average pooling module.

Create a beta schedule that discretizes the given alpha_t_bar function, which defines the cumulative product of (1-beta) over time from t = [

Compute the entire variational lower-bound, measured in bits-per-dim, as well as other related quantities. :param model: the

Compute the mean for the previous step, given a function cond_fn that computes the gradient of a conditional log probability with res

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

Sample x_{t+1} from the model using DDIM reverse ODE.

Sample x_{t-1} from the model using DDIM. Same usage as p_sample().

Compute the log-likelihood of a Gaussian distribution discretizing to a given image. :param x: the target images. It is assumed that thi