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.

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.

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

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

Create a linear module.

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

Zero out the parameters of a module and return it.

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 pas

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

Convert the torso of the model to float16.

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

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

Set the height and weight before and after cropping

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

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

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

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

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

Defaults for image and classifier training.

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

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

Log a value of some diagnostic Call this once for each diagnostic quantity, each iteration If called many times, last value will be used.

Copy the master parameter data back into the model parameters.

Defaults for image training.

Copy the gradients from the model parameters into the master parameters from make_master_params().

Copied from: https://github.com/openai/baselines/blob/ea25b9e8b234e6ee1bca43083f8f3cf974143998/baselines/common/mpi_util.py#L110 Perform a we

Sample x_{t-1} from the model at the given timestep. :param model: the model to sample from. :param x: the current tensor at

Generate samples from the model. :param model: the model module. :param shape: the shape of the samples, (N, C, H, W).

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

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