***

layout: overview
slug: nemo-curator/nemo\_curator/models/cosmos\_embed1
title: nemo\_curator.models.cosmos\_embed1
------------------------------------------

## Module Contents

### Classes

| Name                                                              | Description                    |
| ----------------------------------------------------------------- | ------------------------------ |
| [`CosmosEmbed1`](#nemo_curator-models-cosmos_embed1-CosmosEmbed1) | Cosmos-Embed1 embedding model. |

### Data

[`COSMOS_EMBED1_MODEL_REVISION_INFO`](#nemo_curator-models-cosmos_embed1-COSMOS_EMBED1_MODEL_REVISION_INFO)

[`_COSMOS_EMBED1_VARIANTS_INFO`](#nemo_curator-models-cosmos_embed1-_COSMOS_EMBED1_VARIANTS_INFO)

### API

<Anchor id="nemo_curator-models-cosmos_embed1-CosmosEmbed1">
  <CodeBlock showLineNumbers={false} wordWrap={true}>
    ```python
    class nemo_curator.models.cosmos_embed1.CosmosEmbed1(
        variant: typing.Literal['224p', '336p', '448p'] = '336p',
        utils_only: bool = False,
        model_dir: str | None = None
    )
    ```
  </CodeBlock>
</Anchor>

<Indent>
  **Bases:** [ModelInterface](/nemo-curator/nemo_curator/models/base#nemo_curator-models-base-ModelInterface)

  Cosmos-Embed1 embedding model.

  <ParamField path="_model" type="AutoModel | None = None" />

  <ParamField path="_weights_dir" type="= str(Path(model_dir) / self._weights_name)" />

  <ParamField path="_weights_name" type="= _COSMOS_EMBED1_VARIANTS_INFO[variant]" />

  <ParamField path="model_id_names" type="list[str]">
    Get the model ID names.
  </ParamField>

  <Anchor id="nemo_curator-models-cosmos_embed1-CosmosEmbed1-download_processor_config_on_node">
    <CodeBlock showLineNumbers={false} wordWrap={true}>
      ```python
      nemo_curator.models.cosmos_embed1.CosmosEmbed1.download_processor_config_on_node(
          model_dir: str,
          variant: typing.Literal['224p', '336p', '448p'] = '336p'
      ) -> None
      ```
    </CodeBlock>
  </Anchor>

  <Indent>
    <Badge>
      classmethod
    </Badge>

    Download the processor config for the CosmosEmbed1 model on the node.
  </Indent>

  <Anchor id="nemo_curator-models-cosmos_embed1-CosmosEmbed1-download_weights_on_node">
    <CodeBlock showLineNumbers={false} wordWrap={true}>
      ```python
      nemo_curator.models.cosmos_embed1.CosmosEmbed1.download_weights_on_node(
          model_dir: str,
          variant: typing.Literal['224p', '336p', '448p'] = '336p'
      ) -> None
      ```
    </CodeBlock>
  </Anchor>

  <Indent>
    <Badge>
      classmethod
    </Badge>

    Download the weights for the CosmosEmbed1 model on the node.
  </Indent>

  <Anchor id="nemo_curator-models-cosmos_embed1-CosmosEmbed1-encode_video_frames">
    <CodeBlock showLineNumbers={false} wordWrap={true}>
      ```python
      nemo_curator.models.cosmos_embed1.CosmosEmbed1.encode_video_frames(
          frames: numpy.typing.NDArray[numpy.float32]
      ) -> torch.Tensor
      ```
    </CodeBlock>
  </Anchor>

  <Indent>
    Encode video frames for the model.

    **Parameters:**

    <ParamField path="frames" type="npt.NDArray[np.float32]">
      The input video frames.
    </ParamField>

    **Returns:** `torch.Tensor`

    The encoded video frames.
  </Indent>

  <Anchor id="nemo_curator-models-cosmos_embed1-CosmosEmbed1-evaluate">
    <CodeBlock showLineNumbers={false} wordWrap={true}>
      ```python
      nemo_curator.models.cosmos_embed1.CosmosEmbed1.evaluate(
          video_embd: torch.Tensor,
          text_embds: list[torch.Tensor]
      ) -> tuple[list[float], list[int]]
      ```
    </CodeBlock>
  </Anchor>

  <Indent>
    Evaluate the model.

    **Parameters:**

    <ParamField path="video_embd" type="torch.Tensor">
      The video embedding.
    </ParamField>

    <ParamField path="text_embds" type="list[torch.Tensor]">
      The text embeddings.
    </ParamField>

    **Returns:** `tuple[list[float], list[int]]`

    The predicted probabilities and indices.
  </Indent>

  <Anchor id="nemo_curator-models-cosmos_embed1-CosmosEmbed1-formulate_input_frames">
    <CodeBlock showLineNumbers={false} wordWrap={true}>
      ```python
      nemo_curator.models.cosmos_embed1.CosmosEmbed1.formulate_input_frames(
          frames: list[numpy.typing.NDArray[numpy.uint8]]
      ) -> numpy.typing.NDArray[numpy.float32] | None
      ```
    </CodeBlock>
  </Anchor>

  <Indent>
    Formulate input frames for the model.

    **Parameters:**

    <ParamField path="frames" type="list[npt.NDArray[np.uint8]]">
      List of video frames.
    </ParamField>

    **Returns:** `npt.NDArray[np.float32] | None`

    The formulated input frames.
  </Indent>

  <Anchor id="nemo_curator-models-cosmos_embed1-CosmosEmbed1-get_target_num_frames">
    <CodeBlock showLineNumbers={false} wordWrap={true}>
      ```python
      nemo_curator.models.cosmos_embed1.CosmosEmbed1.get_target_num_frames() -> int
      ```
    </CodeBlock>
  </Anchor>

  <Indent>
    Get the target number of frames for the model.

    **Returns:** `int`

    The target number of frames.
  </Indent>

  <Anchor id="nemo_curator-models-cosmos_embed1-CosmosEmbed1-get_text_embedding">
    <CodeBlock showLineNumbers={false} wordWrap={true}>
      ```python
      nemo_curator.models.cosmos_embed1.CosmosEmbed1.get_text_embedding(
          text: str
      ) -> torch.Tensor
      ```
    </CodeBlock>
  </Anchor>

  <Indent>
    Get the text embedding for the given text.

    **Parameters:**

    <ParamField path="text" type="str">
      The input text.
    </ParamField>

    **Returns:** `torch.Tensor`

    The text embedding.
  </Indent>

  <Anchor id="nemo_curator-models-cosmos_embed1-CosmosEmbed1-setup">
    <CodeBlock showLineNumbers={false} wordWrap={true}>
      ```python
      nemo_curator.models.cosmos_embed1.CosmosEmbed1.setup() -> None
      ```
    </CodeBlock>
  </Anchor>

  <Indent>
    Set up the Cosmos-Embed1 model.

    This method initializes the model and its configuration for processing video and text data.
  </Indent>
</Indent>

<Anchor id="nemo_curator-models-cosmos_embed1-COSMOS_EMBED1_MODEL_REVISION_INFO">
  <CodeBlock showLineNumbers={false} wordWrap={true}>
    ```python
    nemo_curator.models.cosmos_embed1.COSMOS_EMBED1_MODEL_REVISION_INFO: Final = {'224p': '85f5627', '336p': '5d8309d', '448p': '9f4ff4d'}
    ```
  </CodeBlock>
</Anchor>

<Anchor id="nemo_curator-models-cosmos_embed1-_COSMOS_EMBED1_VARIANTS_INFO">
  <CodeBlock showLineNumbers={false} wordWrap={true}>
    ```python
    nemo_curator.models.cosmos_embed1._COSMOS_EMBED1_VARIANTS_INFO: Final = {'224p': 'nvidia/Cosmos-Embed1-224p', '336p': 'nvidia/Cosmos-Embed1-336p', '448p...
    ```
  </CodeBlock>
</Anchor>