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

Take the mean over all non-batch dimensions.

Input: points: input points data, [B, N, C] idx: sample index data, [B, S] Return: new_points:, indexed points data,

Appends dimensions to the end of a tensor until it has target_dims dimensions.

Embed images into latent grids. :param xs: an iterable of images to embed. :return: a tensor of shape [N x C x L], where L =

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

Converts a denoiser output to a Karras ODE derivative.

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 allow

:param xs: N images, stored as numpy arrays, tensors, or PIL images. :return: an [N x D] tensor of features.

Embed text prompts as an [N x D] tensor.

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

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

Get the device that should be used for input tensors.

Encode a batch of point clouds to cache part of the SDF calculation done by forward(). :param point_clouds: a batch of [batc

Input: xyz: pointcloud data, [B, N, 3] npoint: number of samples Return: centroids: sampled pointcloud index, [B, npo

Extract the info of every array in an npz file.

Load the mesh from a .npz file.

For each point in another set of points, compute the point in this pointcloud which is closest. :param points: an [N x 3] ar

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

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

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

Input: radius: local region radius nsample: max sample number in local region xyz: all points, [B, N, 3] new_xyz:

Calculate Euclid distance between each two points. src^T * dst = xn * xm + yn * ym + zn * zm; sum(src^2, dim=-1) = xn*xn + yn*yn + zn*zn

Write a PLY file for a mesh or a point cloud. :param coords: an [N x 3] array of floating point coordinates. :param rgb: an [N x 3] arra

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

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

Like forward(), but always runs all encoders to ensure that the forward graph looks the same on every rank.

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

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_pro

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

Download the file at the given URL into a local file and return the path. If cache_dir is specified, it will be used to download the files.

Compute the Frechet distance between two sets of statistics.

Calculates the noise level (sigma_down) to step down to and the amount of noise to add (sigma_up) when doing an ancestral sampling step.

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

Get a pre-defined beta schedule for the given name. The beta schedule library consists of beta schedules which remain similar in the lim

Constructs the noise schedule of Karras et al. (2022).