Convert rotation matrix/matrices to quaternion(s). Args: R (torch.Tensor [...,3,3]): Rotation matrix/matrices. eps (f

Generates a video based on an input image or video and text prompt. Processes the input, prepares the data batch, runs the diffusion

Get next batch of actions. Returns shape (num_actions, action_dim).

This function creates buffer name used by EMA from parameter's name Args: param_name (str): Model's parameter name Returns:

Get the camera center in world coordinates for a given world2cam pose. Args: cam_pose (torch.Tensor/np.ndarray [...,3,4]): Camera

Return checkpoint config for S3 URI.

Get the metadata for a key in the LeRobot modality metadata. Args: key (str): The key to get the metadata for. Returns:

Get model literal for a given variant.

Get Plücker coordinates (moment, direction) for each pixel center. Args: cam_pose (torch.Tensor/np.ndarray [...,3,4]): Camera pos

Get the results of previously submitted checkpoints and pass them to callbacks if checkpoint succeeded

Generate a sequence of reverse time steps. Args: t_min (float): The minimum time value. t_max (float): The maximum time valu

Calculate and return statistics across all metrics. Returns: Dictionary containing averaged metrics

Generates a callable function `x0_fn` based on the provided data batch and guidance factor. This function first processes the input

Import a module by name. Args: full_module_name: The fully qualified name of the module to import. reload: If True, reload t

Initialize environment.

Initialize output directory.

Extract (fx, fy, cx, cy) from camera intrinsic matrix/matrices. Args: cam_intr (torch.Tensor/np.ndarray [...,3,3]): Camera intrin

Check if current process is the master GPU. Returns: (bool): True if this function is called from the master GPU, else False.

List all keys in an S3 bucket with a given prefix, excluding files that start with specified prefix. Args: filep

Scan a directory to find the interested directories or files in arbitrary order. Note: Boto3 has no concept of directorie

Hamilton product of two quaternion sets. Args: q1 (torch.Tensor [...,4]): Left quaternion(s) in (x, y, z, w) order. q

Hook before zero_grad() is called. Args: optimizer (torch.optim.Optimizer): The model optimizer. scheduler (torch.opt

Apply tensor parallelism, activation checkpointing, torch.compile, and data parallelism to the model. NOTE: The passed-in model preferab

Read keyword list from all files in a folder.

Perform a regularized Euler step based on epsilon prediction. Args: x_s: Current state tensor. s: Current time tensor.

Register a backend. Args: name (str): The name of the registered backend. backend (class, optional): The backend class to be regi

Reset all metrics to their default values.

Rotates half the hidden dims of the input.

Run VQA evaluation on all MP4 files in the configured directory. Returns: Dictionary mapping video filenames to their re

Save network weights, optimizer parameters, scheduler parameters to a checkpoint. Args: model (ImaginaireModel): The PyTorch mode

Save network weights, optimizer parameters, scheduler parameters to a checkpoint. Args: model_parts (list[nn.Module]): The PyTorc

Helper function for saving object to storage. Args: object (Any): The object to upload. key (str): The key of the obj

Set chunk size Args: chunk_size (int): chunk size used in webdataset creation

Set urls Args: urls (list[TarSample]): a list of tar files along with their metadata

Convert sRGB values to physically linear ones. The transformation is uniform in RGB, so *srgb* can be of any shape. *srgb* values should rang

Get homogeneous coordinates of the input points. Args: points (torch.Tensor/np.ndarray [...,K]): Input coordinates. Retu

Convert tensor data type from uint8 to float, divide value by 255.0 and permute the dimensions of clip tensor Args: clip (torch.t

The training function. Args: model (ImaginaireModel): The PyTorch model. dataloader_train (torch.utils.data.DataLoade

Initialize the dataset. Args: dataset_path (Path | str): The path to the dataset. modality_configs (dict[str

Initialize S3FileSystem with retry configuration. Args: credential_path: Path to AWS credentials JSON file m

Initialize the depth to point cloud converter. Args: input_keys: List of input keys (typically ['depth', 'intrinsics', 'world_to_

Inverse of ``locate()``. Args: t: any object with ``__module__`` and ``__qualname__``

r"""This function performs tar file equalization. That is, we repeat the number of tars in each aspect ratio so that when the tars are split a

Processes the input batch through all configured embedders, applying conditional dropout rates if specified. Output tensors for each

Get the interpolation mode for the given backend. Args: interpolation (str): The interpolation mode. backend

Get the rank (GPU device) of the worker. Returns: rank (int): The rank of the worker.

Args: text (str): string in which we will search for the word word_ind (int): index of the word inside of this string

Initialize the S3 client.

Lazy-load the Cosmos Reason model.