Forward pass of the model. Args: x: input motion, [bs, seq_len, motion_dim] kwargs: other inputs, e.g., contact, tex

... :param x_q: Query input of shape (B, N, D) where B is the batch size, N the query sequence length and D the number of que

Pre-layer-norm cross-attention of query input `x_q` to key/value input (`x_kv` or `x_kv_prefix`). If `x_kv_prefix` is defined, the entire key

Pre-layer-norm self-attention of input `x`.

Forward pass of the ContactMLP. Args: x: input contact map, [bs, num_points, contact_dim] point_feat: [bs, num_point

Forward pass of the ContactMLP. Args: x: input contact map, [bs, num_points, contact_dim] point_feat: [bs, num_point

Forward pass of the ContactMLP. Args: x: input contact map, [bs, num_points, contact_dim] point_feat: [bs, num_point

Forward pass of the ContactMLP. Args: x: input contact map, [bs, num_points, contact_dim] point_feat: [bs, num_point

Forward pass of the model. Args: x: input contact map, [bs, num_points, contact_dim] kwargs: other inputs, e

input: xyz: (n, 3), offset: (b), new_offset: (b) output: idx: (m)

input: xyz: (n, 3), new_xyz: (m, 3), offset: (b), new_offset: (b) output: idx: (m, nsample), dist2: (m, nsample)

input: input: (n, c), idx : (m, nsample) output: (m, nsample, c)

input: input1: (n, c), input2: (n, c), idx: (n, nsample) output: (n, nsample, c)

input: input: (n, c), position: (n, nsample, c), weight : (n, nsample, c'), idx: (n, nsample) output: (n, c)

input: xyz: (m, 3), new_xyz: (n, 3), input: (m, c), offset: (b), new_offset: (b) output: (n, c)

Get dataloader

Get dataloader

Get dataloader

Get dataloader

Get dataloader

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

input: xyz: (m, 3), new_xyz: (n, 3), feat: (m, c), offset: (b), new_offset: (b) output: (n, c)

Create a linear module.

Load humanise dataset Args: min_horizon: minimum horizon for motion sequence max_horizon: maximum horizon for motion sequence

Process humanml3d dataset Args: min_horizon: minimum horizon for motion sequence max_horizon: maximum horizon for motion sequenc

Load prox dataset Args: min_horizon: minimum horizon for motion sequence max_horizon: maximum horizon for motion sequence

Make a standard normalization layer. :param channels: number of input channels. :return: an nn.Module for normalization.

Generate samples from the model. :param model: the model module. :param shape: the shape of the samples, (N, C, H, W).

Pad a batch motion sequence to the same length

Convert point cloud to cubes Args: xyz: point cloud, <N, 3> Return: Cube mesh list

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

input: xyz: (n, 3), new_xyz: (m, 3), feat: (n, c), idx: (m, nsample), offset: (b), new_offset: (b) output: new_feat: (m, c+3, nsample), group

Generate a random string with the given length Args: length: length of the random string Return: random string

Register the given object under the the name `obj.__name__`. Can be used as either a decorator or not. See docstring of this class fo

Render mesh with contact map Args: save_path: save path mesh: mesh to render camera_pose: camera pose

Report evaluation results Args: save_dir: save directory

Report evaluation results Args: save_dir: save directory

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

Scale the parameters of a module and return it.

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

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

Update target parameters to be closer to those of source parameters using an exponential moving average. :param target_params: the targe

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

Visualize samples Args: samples: samples to visualize save_dir: directory to save samples

Visualize samples Args: save_dir: directory to save samples samples: samples to visualize

Visualize samples Args: save_dir: directory to save samples samples: samples to visualize

Visualize contact map on scene point cloud Args: xyz: scene point cloud xyz rgb: scene point cloud rgb contact: cont

Visualize partial scene and skeleton Args: partial_scene: partial scene point cloud pose_seq: pose sequence scene_tr

Visualize reconstruction results Args: x: grount truth data rec_x: reconstruction data save_dir: directo