A dataset that generates items by applying a function. This allows for creating dynamic datasets where the items are the result of function c

Config for checkpoint file on Hugging Face.

Config for checkpoint.

Config for checkpoint directory on S3.

Configuration for a modality.

Wrapper for model state dict handling

Causal 3d convolusion.

Causal 3d convolusion.

r""" This class is a wrapper on webloader class with a len attribute. len function is needed in Imaginaire dataloaders.

Config for checkpoint directory on Hugging Face.

Config for checkpoint file on S3.

Similar as FastEmaModelUpdater

This class is used to update target model~(EMA) given source model~(regular model) and beta. The method interaface mimic :class:`EMAModelTrac

Subclass torch's LayerNorm to handle fp16.

Wrapper for model state dict handling

Convert tensor data type from uint8 to float, divide value by 255.0 and permute the dimensions of clip tensor

Causal self-attention with a single head.

Class used to store times computed during tokenizer benchmarking. All times are in seconds.

A dataset that wraps multiple PyTorch datasets and returns a dictionary of data items from each dataset for a given index. This dataset ensur

Config for the EMA.

Wrapper for model state dict handling

RenameLoadPlanner that renames variables during checkpoint load.

Video preprocessing class for resizing frames to a specific size.

Transforms states and actions to tensors.

Class for state or action transform. Args: apply_to (list[str]): The keys in the modality to load and transform. normalizati

Config for the text encoder model

Handles the Video2World inference process, including model loading, data preparation, and video generation from an image/video and text promp

A flexible attention module supporting both self-attention and cross-attention mechanisms. This module implements a multi-head attention lay

A flexible attention module supporting both self-attention and cross-attention mechanisms. This module implements a multi-head attention lay

A DataLoader wrapper that asynchronously caches and replays data batches to mitigate slow loading issues. Assumes the underlying DataLoader is inf

Compose multiple modality transforms.

`UniPCMultistepScheduler` is a training-free framework designed for the fast sampling of diffusion models. This model inherits from [`Schedu

Evaluates VQA checks on a single video using Cosmos Reason model. This class manages the model lifecycle and provides methods to run VQA eva

Represents the result of a VQA check. Attributes: check: The original VQACheck object actual_answer: The actual answer gener

r"""Webdataset class modified to support loading from object store.

Records and tracks various loss metrics during training.

Simple CPU timer for timing the execution of code. It can be used as either a context manager or a function decorator. The timing result will be

A flexible attention module supporting both self-attention and cross-attention mechanisms. This module implements a multi-head attention lay

Causal self-attention with a single head.

Logger that does nothing, used when logging is disabled.

Bottleneck residual block for ResNet.

Checkpointer class for Tensor Parallelism (TP) in distributed training. This implementation supports the combination of Tensor Parallelism (

Concatenate the keys according to specified order.

Handles T5 text encoding operations.

A dummy dataset that returns a dictionary with a 'videos' tensor of shape (1, 3, frames_per_batch). Each item is a tensor with sequential values

This is a class to track the EMA model weights. The EMA weights are registered as buffers, which are extractable as state dicts. The names follow