Autoencoder3d adapted from Stable Diffusion 1.x, 2.x and XL.

Determine dataset type from file path. Args: path: Path to the file dataset_names: List of dataset names to check agains

Initialize the step_index counter for the scheduler.

Initialize the step_index counter for the scheduler.

Process frame attention blocks. We keep tokens in shape (B*S, P, C).

Process global attention blocks. We keep tokens in shape (B, S*P, C).

Performs feature rotation by splitting and recombining feature dimensions. Args: x: Input tensor to rotate. Returns:

Returns torch.sqrt(torch.max(0, x)) but with a zero subgradient where x is 0.

Autoencoder3d adapted from Stable Diffusion 1.x, 2.x and XL.

Process network output to extract 3D points and confidence values. Args: out: Network output tensor (B, C, H, W) activation:

Activate pose parameters with specified activation functions. Args: pred_pose_enc: Tensor containing encoded pose parameters [transl

Computes translation, rotation, and focal loss for a batch of pose encodings. Args: pred_pose_enc: (N, D) predicted pose encodin

Compute the inverse of each 4x4 (or 3x4) SE3 matrix in a batch. If `R` and `T` are provided, they must correspond to the rotation and transl

Transfer color distribution from of sc, referred to dc. Args: sc (numpy.ndarray): input image to be transfered. dc (numpy.nd

Convert the model output to the corresponding type the DPMSolver/DPMSolver++ algorithm needs. DPM-Solver is designed to discretize an

r""" Convert the model output to the corresponding type the UniPC algorithm needs. Args: model_output (`torch.Tensor`):

Helper function to create dataset configs.

Create a normalized UV grid of shape (width, height, 2). The grid spans horizontally and vertically according to an aspect ratio, ensuri

returns either a context list that includes one that will disable zero.Init or an empty context list

Convert a depth map to camera coordinates. Args: depth_map (np.ndarray): Depth map of shape (H, W). intrinsic (np.ndarray):

Convert a depth map to world coordinates. Args: depth_map (np.ndarray): Depth map of shape (H, W). intrinsic (np.ndarray): C

One step for the first-order DPMSolver (equivalent to DDIM). Args: model_output (`torch.Tensor`): The dir

r""" Encodes the prompt into text encoder hidden states. Args: prompt (`str` or `List[str]`, *optional*):

r""" Args: x(Tensor): Shape [B, L1, C] e(Tensor): Shape [B, C]

x_list contains a list of tensors to nest together and run

RIFLEx: Precompute the frequency tensor for complex exponentials (cis) with given dimensions. This function calculates a frequency tensor wi

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

Modulate the input tensor using scaling and shifting parameters.

One step for the second-order multistep DPMSolver. Args: model_output_list (`List[torch.Tensor]`): The di

One step for the third-order multistep DPMSolver. Args: model_output_list (`List[torch.Tensor]`): The dir

One step for the UniC (B(h) version). Args: this_model_output (`torch.Tensor`): The model outputs at `x_

One step for the UniP (B(h) version). Alternatively, `self.solver_p` is used if is specified. Args: model_output (`torch

Convert 2D position grid (HxWx2) to sinusoidal embeddings (HxWxC) Args: pos_grid: Tensor of shape (H, W, 2) containing 2D coordinate

Loads and preprocesses a video. Args: video_path: Path to the video file to load. Returns: torch.Te

Loads and preprocesses a camera path.

Quaternion Order: XYZW or say ijkr, scalar-last Convert rotations given as quaternions to rotation matrices. Args: quaternions:

Ensures interchangeability with schedulers that need to scale the denoising model input depending on the current timestep. Ar

Forward pass through the fusion blocks. Args: features (List[Tensor]): List of feature maps from different layers.

Select prompts for validation with one fixed sample and the rest randomly sampled. Selection strategy: 1) Always select one prompt fro

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

Predict the sample from the previous timestep by reversing the SDE. This function propagates the sample with the multistep UniPC.

Better torch.quantile for one SCALAR quantile. Using torch.kthvalue. Better than torch.quantile because: - No 2**24 input size limit (pyt

Iteratively refine camera pose predictions. Args: pose_tokens (torch.Tensor): Normalized camera tokens with shape [B, 1,

r""" Reconstruct video tensors from patch embeddings. Args: x (List[Tensor]): List of patchified features

Applies transformations to a video. Args: frames (torch.Tensor): A 4D tensor representing a video with s