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

Setup distributed training if script is launched with torchrun. Currently only supports single node multi GPU.

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

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

Flattens predictions in the batch DynamicCity: changed some dimensions

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

Padded rollout, converts hexplane to (b, c, x+t+z, y+z+t) - xy: :x, :y - xz: :x, y:y+z - zy: x+t:, - tx: :x, y+z: - ty: x:x+t

Multi-class Lovasz-Softmax loss probas: [P, C] Variable, class probabilities at each prediction (between 0 and 1) labels: [P] Tensor,

nanmean compatible with generators.

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

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

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

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

model_out should contain 'pred' for ce and lovasz, 'mus' and 'logvars' for kl

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

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

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

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