Collect results in cpu.

Collect results in gpu.

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 co

Convert applicable model parameters to fp32

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

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

Generate samples from the model using DDIM. Same usage as p_sample_loop().

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

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

Get target label. Args: input (Tensor): Input tensor. target_is_real (bool): Whether the target is real or fake.

Geman-McClure error function.

Load annotations from ``ann_file`` to ``data_infos``

Convert rotations given as rotation matrices to quaternions. Args: matrix: Rotation matrices as tensor of shape (..., 3, 3). Re

Warpper of mse loss.

Test model with multiple gpus. This method tests model with multiple gpus and collects the results under two different modes: gpu and cpu mode

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_

Parse a version string into a tuple. Args: version_str (str): The version string. Returns: tuple[int | str]: The version info,

Prepare raw data for the f'{idx'}-th data.

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 dif

find the quaternion used to rotate v0 to v1

find the quaternion used to rotate v0 to v1

q0 : tensor of quaternions t: tensor of powers

Rotate vector(s) v about the rotation described by quaternion(s) q. Expects a tensor of shape (*, 4) for q and a tensor of shape (*, 3) for v

Given a quaternion representing rotation, get the quaternion representing its inverse. Args: quaternion: Quaternions as tensor o

Convert rotations given as quaternions to axis/angle. Args: quaternions: quaternions with real part first, as tensor of

Generate random quaternions representing rotations, i.e. versors with nonnegative real part. Args: n: Number of quaternions in a

Generate random rotations as 3x3 rotation matrices. Args: n: Number of rotation matrices in a batch to return. dtype: Type t

Reduce loss as specified. Args: loss (Tensor): Elementwise loss tensor. reduction (str): Options are "none", "mean" and "sum".

Set random seed. Args: seed (int): Seed to be used. deterministic (bool): Whether to set the deterministic option for

Test with single gpu.

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

Convert a unit quaternion to a standard form: one in which the real part is non negative. Args: quaternions: Quaternions with re

Wrap models with separate MMDistributedDataParallel. It only wraps the modules with parameters.

Convert objects of various python types to :obj:`torch.Tensor`. Supported types are: :class:`numpy.ndarray`, :class:`torch.Tensor`, :class:`S

Main api for training model.

Compute training losses for a single timestep. :param model: the model to evaluate loss on. :param x_start: the [N x C x ...

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

Update the reweighting using losses from a model. Call this method from each rank with a batch of timesteps and the correspon

Apply element-wise weight and reduce loss. Args: loss (Tensor): Element-wise loss. weight (Tensor): Element-wise weights.

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

Zero out the parameters of a module and return it.

Prepare data for the ``idx``-th data. As for video dataset, we can first parse raw data for each frame. Then we combine annotations fr