Standard CutMix bounding-box Generates a random square bbox based on lambda value. This impl includes support for enforcing a border margin a

Min-Max CutMix bounding-box Inspired by Darknet cutmix impl, generates a random rectangular bbox based on min/max percent values applied to e

Removes the redundant lines added during cropping.

Remove instances that have an area ratio smaller than `bbox_area_threshold`.

Render OpenPose body keypoints on input image. Args: img (np.array): Input image of shape (H, W, 3) with pixel values in the [0,255]

Render keypoints on input image. Args: img (np.array): Input image of shape (H, W, 3) with pixel values in the [0,255] range.

Rescale `noise_cfg` according to `guidance_rescale`. Based on findings of [Common Diffusion Noise Schedules and Sample Steps are Flawed](http

Reset the iterator when exhausted.

Resize the sample

Returns a mask of shape B H_Q W_Q, where True = masked-out, False = keep. Args: quant: Dequantized latent tensor of shape B D_Q H

Returns a mask of shape B H_Q W_Q, where True = masked-out, False = keep. Args: quant: Dequantized latent tensor of shape B D_Q H

Saves the conditioning and target modalities for a batch of samples. Args: cond_domains (list of str): List of conditioning domains

Applies input masking to an sequence embedding tensor, target masking is not supported with sequence embeddings Args: emb_tensor:

Applies input and target masking to a sequence The keep probability is sampled from a cosine schedule and does not depend on the number of to

Applies input and target masking to a sequence of tokens (e.g. DINOv2 global tokens) The keep probability is sampled from a cosine schedule an

Sets the discrete timesteps used for the diffusion chain. Supporting function to be run before inference. Args: num_infe

This function disables printing when not in main process

Sets up the modality info dictionary for the given domains.

Sets up the modality info dictionary for the given domains.

Sets up the modality info dictionary for the given domains.

Share modality embeddings across the encoder and decoder embedding modules.

Step function propagating the sample with the linear multi-step method. This has one forward pass with multiple times to approximate

Step function propagating the sample with the Runge-Kutta method. RK takes 4 forward passes to approximate the solution to the differ

Applies masking for an image given as target Args: tensor: Image tensor num_tokens: Number of tokens in the tensor

Applies masking for a sequence given as target Args: sequence_str: Sequence string max_tokens: Maximum number of toke

Sample a token budget for the target Args: input_token_budget: Token budget for the input num_target_tokens: Number o

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

Tokenizes an input image tensor. Args: x: Input image tensor of shape B C H W or B H W in case of semantic segmenta

Look up the codebook entries corresponding the discrete tokens. Args: tokens: Quantized indices of shape B H_Q W_Q Retur

Perform one training epoch and return stats. The image resolution is randomly sampled from train_res_choices. At the specified intervals, eva

Perform one training epoch and return stats. The image resolution is randomly sampled from train_res_choices. At the specified intervals, eva

If separate==True, the transforms are returned as a tuple of 3 separate transforms for use in a mixing dataset that passes * all data th

Truncated depth standardization :param depth: Depth map :param thresh: Threshold :return: Robustly standardized depth map

Converts all `FrozenBatchNorm2d` layers of provided module into `BatchNorm2d`. If `module` is itself and instance of `FrozenBatchNorm2d`, it

Apply the image augmenter to all modalities where it is applicable Args: mod_dict (dict): Dict of modalities crop_set

Args: img (PIL Image): Image to be cropped and resized. Returns: PIL Image: Randomly cropped and resized ima

Render meshes on input image Args: vertices (np.array): Array of shape (V, 3) containing the mesh vertices. c

Render batch of 3D keypoints. Args: pred_keypoints_3d (torch.Tensor): Tensor of shape (B, S, N, 3) containing a batch of

Applies input and target masking to a dictionary of modalities Args: mod_dict: Dictionary of modalities Returns:

Applies input and target masking to a dictionary of modalities Args: mod_dict: Dictionary of modalities Returns:

Apply the augmentation to a dict of modalities (both image based and sequence based modalities) Args: mod_dict (dict): Dict of mo

The call function to the pipeline for conditional image generation. Args: cond: The conditional input to the model. g