`stages.video.embedding.cosmos_embed1`#

Module Contents#

Classes#

`CosmosEmbed1EmbeddingStage`	Stage for embedding Cosmos-Embed1 frames into a vector space.
`CosmosEmbed1FrameCreationStage`	Stage for creating Cosmos-Embed1 input frames from video clips.

API#

class stages.video.embedding.cosmos_embed1.CosmosEmbed1EmbeddingStage#

Bases: nemo_curator.stages.base.ProcessingStage[nemo_curator.tasks.video.VideoTask, nemo_curator.tasks.video.VideoTask]

Stage for embedding Cosmos-Embed1 frames into a vector space.

This class processes video clips through a series of steps including frame extraction, model initialization, and input frame creation.

gpu_memory_gb: int#: 20

inputs() → tuple[list[str], list[str]]#

Define stage input requirements.

Returns (tuple[list[str], list[str]]): Tuple of (required_attributes, required_columns) where: - required_top_level_attributes: List of task attributes that must be present - required_data_attributes: List of attributes within the data that must be present

model_dir: str#: ‘models/cosmos_embed1’

outputs() → tuple[list[str], list[str]]#

Define stage output specification.

Returns (tuple[list[str], list[str]]): Tuple of (output_attributes, output_columns) where: - output_top_level_attributes: List of task attributes this stage adds/modifies - output_data_attributes: List of attributes within the data that this stage adds/modifies

process( task: nemo_curator.tasks.video.VideoTask, ) → nemo_curator.tasks.video.VideoTask#: Process a task and return the result. Args: task (X): Input task to process Returns (Y | list[Y]): - Single task: For 1-to-1 transformations - List of tasks: For 1-to-many transformations (e.g., readers) - None: If the task should be filtered out

setup( worker_metadata: nemo_curator.backends.base.WorkerMetadata | None = None, ) → None#: Setup method called once before processing begins. Override this method to perform any initialization that should happen once per worker. Args: worker_metadata (WorkerMetadata, optional): Information about the worker (provided by some backends)

setup_on_node( node_info: nemo_curator.backends.base.NodeInfo, worker_metadata: nemo_curator.backends.base.WorkerMetadata, ) → None#: Download the weights for the CosmosEmbed1 model on the node.

texts_to_verify: list[str] | None#: None

variant: Literal[224p, 336p, 448p]#: ‘336p’

verbose: bool#: False

class stages.video.embedding.cosmos_embed1.CosmosEmbed1FrameCreationStage#