Load a checkpoint. Should be called by all processes.

Load a pretrained model.

Cast if autocast is not enabled.

Register a spatial cache. The spatial cache can be any thing you want to cache. The registery and retrieval of the cache is b

Get a spatial cache.

Sample from a model.

Save a checkpoint. Should be called only by the rank 0 process.

Convert a state_dict to master params.

Reduce a list of dictionaries. Leaf values must be scalars.

Recursively move all tensors in a data structure to a device.

Render the mesh. Args: mesh : meshmodel extrinsics (torch.Tensor): (4, 4) camera extrinsics intr

Split indices into groups with balanced load.

Read the binary file distributedly. File is only read once by the rank 0 process and broadcasted to other processes. Returns: da

Zero out the parameters of a module and return it.

Check if DDP is working properly. Should be called by all process.

Lazy-load pretrained NAF model.

Render a batch of representations. Args: reps: The dictionary of lists of representations. extrinsics: The [

Run training.

Initialize the image conditioning model.

Render a batch of representations. Args: reps: The dictionary of lists of representations. extrinsics: The [

Diffuse the data for a given number of diffusion steps. In other words, sample from q(x_t | x_0). Args: x_0: The

Encode the image (standard mode without projection).

Get the conditioning data.

Get the conditioning data for inference.

Get the sampler for the diffusion process.

Compute the velocity of the diffusion process at time t.

Override to handle: 1. Old checkpoints that don't have image_cond_step*.pt 2. Partial image_cond checkpoints (only proj_linea

Return a normalization layer.

Sample timesteps.

Sample structured latent with the given conditioning. Args: cond (dict): The conditioning information.

Apply scaled dot product attention. Args: qkv (torch.Tensor): A [N, L, 3, H, C] tensor containing Qs, Ks, and Vs.

Sample images from the model. NOTE: When num_samples >= 4, this function should be called by all processes. When num_sa

Tries to extract an image and a factor from a texture node link.

Visualize the conditioning data.

Extract camera info from kwargs. Supports two formats: 1. 'camera_info' dict: {'camera_angle_x': ..., 'distance': ..

Fill missing per-block proj_linear weights in denoiser checkpoint from the old-style image_cond proj_linear (broadcast to all blocks)

Atomically write progress JSON to a file (cross-process safe).

Calculate serialization and partitioning for a set of coordinates. Args: tensor (SparseTensor): The input tensor. window_siz

Recursively apply a function to all non-dictionary leaf values in a dictionary.

Encode the text.

Create a VarLenTensor from a list of tensors.

Compute camera-to-world 4x4 transform matrix in spherical coordinates, looking at origin with up=(0,1,0). Uses standard LookAt formula wh