Implementation:NVIDIA NeMo Curator CosmosEmbed1EmbeddingStage

Knowledge Sources	NeMo Curator
Domains	Data_Curation, Video_Processing, Representation_Learning
Last Updated	2026-02-14 17:00 GMT

Overview

Concrete tool for computing video clip embeddings using Cosmos-Embed1 provided by NeMo Curator.

Description

The CosmosEmbed1EmbeddingStage computes dense vector embeddings for video clips using the Cosmos-Embed1 model. It supports multiple resolution variants (224p, 336p, 448p) and optional text-video similarity verification. Preceded by CosmosEmbed1FrameCreationStage which prepares model input frames.

Usage

Import this stage after CosmosEmbed1FrameCreationStage to compute clip embeddings for semantic deduplication.

Code Reference

Source Location

• Repository: NeMo Curator
• File: nemo_curator/stages/video/embedding/cosmos_embed1.py
• Lines: L109-167

Signature

@dataclass
class CosmosEmbed1EmbeddingStage(ProcessingStage[VideoTask, VideoTask]):
    model_dir: str = "models/cosmos_embed1"
    variant: Literal["224p", "336p", "448p"] = "336p"
    texts_to_verify: list[str] | None = None
    gpu_memory_gb: int = 20
    verbose: bool = False
    name: str = "cosmos_embed1_embedding"

Import

from nemo_curator.stages.video.embedding.cosmos_embed1 import CosmosEmbed1EmbeddingStage, CosmosEmbed1FrameCreationStage

I/O Contract

Inputs

Name	Type	Required	Description
task	VideoTask	Yes	Video with clips having prepared Cosmos-Embed1 input frames

Outputs

Name	Type	Description
task	VideoTask	Video with clip.cosmos_embed1_embedding numpy arrays

Usage Examples

from nemo_curator.stages.video.embedding.cosmos_embed1 import (
    CosmosEmbed1FrameCreationStage,
    CosmosEmbed1EmbeddingStage,
)

# 1. Prepare frames for embedding model
frame_stage = CosmosEmbed1FrameCreationStage(
    model_dir="models/cosmos_embed1", variant="336p", target_fps=2.0,
)

# 2. Compute embeddings
embed_stage = CosmosEmbed1EmbeddingStage(
    model_dir="models/cosmos_embed1", variant="336p", gpu_memory_gb=20,
)

Related Pages

Implements Principle

• Principle:NVIDIA_NeMo_Curator_Video_Embedding
• Environment:NVIDIA_NeMo_Curator_Python_Linux_Base
• Environment:NVIDIA_NeMo_Curator_Video_Codec_Stack
• Environment:NVIDIA_NeMo_Curator_Ray_Cluster
• Heuristic:NVIDIA_NeMo_Curator_GPU_Memory_Resource_Allocation
• Heuristic:NVIDIA_NeMo_Curator_Video_Frame_Sampling_Strategy