Create a 1D, 2D, or 3D convolution module.

Convert primitive modules to float16.

Convert primitive modules to float32, undoing convert_module_to_f16().

Sample x_{t+1} from the model using DDIM reverse ODE.

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().

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

Get the device to use for torch.distributed.

Write all of the diagnostics from the current iteration

Convert rotations given as Euler angles in radians to rotation matrices. Args: euler_angles: Euler angles in radians as tensor of sh

Convert Euler angles to quaternions.

Convert axis-angle rotations (aka exponential maps) to quaternions. Stable formula from "Practical Parameterization of Rotations Using the Ex

x: [batch_size, njoints, nfeats, max_frames], denoted x_t in the paper timesteps: [batch_size] (int)

Returns the angle prior loss for the given pose Args: pose: (Bx[23 + 1] * 3) torch tensor with the axis-angle repres

query (batch, query_dim) key (batch, seq_len, key_dim)

Get directory that log files are being written to. will be None if there is no output directory (i.e., if you didn't call start)

Returns the mean of the mixture

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

Create a linear module.