stages.image.embedders.clip_embedder

Module Contents

Classes

ImageEmbeddingStage

Stage for generating image embeddings using CLIP model.

API

class stages.image.embedders.clip_embedder.ImageEmbeddingStage

Bases: nemo_curator.stages.base.ProcessingStage[nemo_curator.tasks.ImageBatch, nemo_curator.tasks.ImageBatch]

Stage for generating image embeddings using CLIP model.

This class processes image batches through a CLIP model to generate embeddings for each image. It assumes image data is already loaded in ImageObject.image_data and stores embeddings in ImageObject.embedding.

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

None

model_inference_batch_size: int

32

num_gpus_per_worker: float

0.25

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.ImageBatch,

) → nemo_curator.tasks.ImageBatch

Process an image batch to generate embeddings.

Args: task: ImageBatch containing list of ImageObject instances with pre-loaded image_data

Returns: ImageBatch with embeddings stored in ImageObject.embedding

remove_image_data: bool

False

setup(

_worker_metadata: nemo_curator.backends.base.WorkerMetadata | None = None,

) → None

Initialize the CLIP image embedding model.

setup_on_node(

node_info: nemo_curator.backends.base.NodeInfo,

worker_metadata: nemo_curator.backends.base.WorkerMetadata,

) → None

Download the weights for the CLIP model on the node.

verbose: bool

False

yield_next_batch(

task: nemo_curator.tasks.ImageBatch,

) → collections.abc.Generator[nemo_curator.tasks.ImageBatch, None, None]

Yield batches of images from the task.

Args: task: ImageBatch containing list of ImageObject instances with pre-loaded image_data

Yields: Generator[dict[str, torch.Tensor]]: A generator of model inputs for the next batch.