Curate VideoTutorialsPipeline Customization

Adding Custom Models

View as Markdown

Integrate a custom model into a NeMo Curator stage. Package model dependencies in the Python environment or image used by the pipeline.

Before You Start

Before you add a model, prepare your development environment:

How to Add a Custom Model

Implement ModelInterface, then provide the model to a stage that uses it.

Review Model Interface

In NeMo Curator, models inherit from nemo_curator.models.base.ModelInterface and must implement model_id_names and setup:

class ModelInterface(abc.ABC):
    """Abstract base class for models used inside stages."""
    @property
    @abc.abstractmethod
    def model_id_names(self) -> list[str]:
        """Return a list of model IDs associated with this model (for example, Hugging Face IDs)."""
    @abc.abstractmethod
    def setup(self) -> None:
        """Set up the model (load weights, allocate resources)."""

Create New Model

The following example defines a minimal model for demonstration.

from typing import Optional
import numpy as np
import numpy.typing as npt
import torch
from nemo_curator.models.base import ModelInterface
WEIGHTS_MODEL_ID = "example/my-model"
class MyCore(torch.nn.Module):
    def __init__(self, resolution: int = 224):
        super().__init__()
        self.resolution = resolution
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        # Initialize your network here
        self.net = torch.nn.Identity().to(self.device)
    @torch.no_grad()
    def __call__(self, x: npt.NDArray[np.float32]) -> torch.Tensor:
        tensor = torch.from_numpy(x).to(self.device).float()
        return self.net(tensor)
class MyModel(ModelInterface):
    def __init__(self, model_dir: str, resolution: int = 224) -> None:
        self.model_dir = model_dir
        self.resolution = resolution
        self._model: Optional[MyCore] = None
    def model_id_names(self) -> list[str]:
        return [WEIGHTS_MODEL_ID]
    def setup(self) -> None:
        # Load weights from self.model_dir/WEIGHTS_MODEL_ID if needed
        self._model = MyCore(self.resolution)
        self._model.eval()

Let’s go through each part of the code piece by piece.

Define the PyTorch Model

WEIGHTS_MODEL_ID = "example/my-model"  # your huggingface (or other) model id
class MyCore(torch.nn.Module):
    def __init__(self, resolution: int = 224):
        super().__init__()
        # Initialize network and load weights from a local path derived from model_dir and WEIGHTS_MODEL_ID

Provide a model ID, such as a Hugging Face identifier, to cache or fetch weights. Your model class can download weights before setup() using its model-loading method.

Implement the Model Interface

class MyModel(ModelInterface):
	...

Your model implements the interface. It defines methods to declare weight identifiers and to initialize the underlying core network.

    def setup(self) -> None:
        self._model = MyCore(self.resolution)
        self._model.eval()

The setup method initializes the underlying MyCore class that performs the model inference.

    def model_id_names(self) -> list[str]:
        return [WEIGHTS_MODEL_ID]

The model_id_names property returns a list of weight IDs. These typically correspond to model repository names but do not have to.

Set GPU requirements in the stage’s resources, such as gpu_memory_gb or gpus. The stage manages GPU allocation through Resources.

Manage model weights

Set model_dir to the weights location. Mount the weights into the container or download them before the stage runs.

Next Steps

Use the model in a custom stage.