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

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

Compute pairwise distances between two batches of feature vectors.

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

Compute image features for downstream evals. :param batches: a iterator over NHWC numpy arrays in [0, 255]. :return: a tuple

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 logger that writes to a log file and stdout.

Downloads a pre-trained DiT model from the web.

Evaluate precision and recall efficiently. :param features_1: [N1 x D] feature vectors for reference batch. :param radii_1:

Compute the Frechet distance between two sets of statistics.

embed_dim: output dimension for each position pos: a list of positions to be encoded: size (M,) out: (M, D)

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

Evaluate pairwise distances between two batches of feature vectors.

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

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

Step the EMA model towards the current model.

Step the EMA model towards the current model.

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

Copied from: https://github.com/numpy/numpy/blob/fb215c76967739268de71aa4bda55dd1b062bc2e/numpy/lib/format.py#L788-L886 Read from file-like

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

Center cropping implementation from ADM. https://github.com/openai/guided-diffusion/blob/8fb3ad9197f16bbc40620447b2742e13458d2831/guided_diff

Center cropping implementation from ADM. https://github.com/openai/guided-diffusion/blob/8fb3ad9197f16bbc40620447b2742e13458d2831/guided_diff

End DDP training.

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

Create a logger that writes to a log file and stdout.

Builds a single .npz file from a folder of .png samples.

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

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

Finds a pre-trained DiT model, downloading it if necessary. Alternatively, loads a model from a local path.

Forward pass of DiT. x: (N, C, H, W) tensor of spatial inputs (images or latent representations of images) t: (N,) tensor of

grid_size: int of the grid height and width return: pos_embed: [grid_size*grid_size, embed_dim] or [1+grid_size*grid_size, embed_dim] (w/

This is the deprecated API for creating beta schedules. See get_named_beta_schedule() for the new library of schedules.

Insert s without opt_

Trains a new DiT model.

Trains a new DiT model.

Run sampling.

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

Get the distribution q(x_t | x_0). :param x_start: the [N x C x ...] tensor of noiseless inputs. :param t: the number of diff

Set requires_grad flag for all parameters in a model.

Set requires_grad flag for all parameters in a model.

Create a list of timesteps to use from an original diffusion process, given the number of timesteps we want to take from equally-sized portio

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

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

Drops labels to enable classifier-free guidance.

Drops labels to enable classifier-free guidance.

x: (N, T, patch_size**2 * C) imgs: (N, H, W, C)

Update the reweighting using losses from a model. Sub-classes should override this method to update the reweighting using los

Get a numpy array of weights, one per diffusion step. The weights needn't be normalized, but must be positive.