Adding Custom Stages

1
from typing import List
2
3
from nemo_curator.stages.base import ProcessingStage
4
from nemo_curator.stages.resources import Resources
5
from nemo_curator.tasks.video import VideoTask
6
7
class MyCustomStage(ProcessingStage[VideoTask, VideoTask]):
8
    """Example stage that reads and writes to the VideoTask."""
9
10
    name = "my_custom_stage"
11
    resources = Resources(cpus=2.0, gpu_memory_gb=8.0)
12
13
    def setup(self, worker_metadata=None) -> None:
14
        # Initialize models or allocate resources here
15
        pass
16
17
    def process(self, task: VideoTask) -> VideoTask | list[VideoTask]:
18
        # Implement your processing and return the modified task (or list of tasks)
19
        return task

1
# You can override resources at construction time using with_()
2
from nemo_curator.stages.resources import Resources
3
4
stage = MyCustomStage().with_(
5
    resources=Resources(cpus=4.0, gpu_memory_gb=16.0)
6
)

1
def setup(self, worker_metadata=None) -> None:
2
    # Load models, warm up caches, etc.
3
    pass

1
def process(self, task: VideoTask) -> VideoTask | list[VideoTask]:
2
    # Process implementation
3
    return task

1
# In Ray Curator, video data lives in VideoTask.data (a Video) which contains Clips.
2
# You can attach new information to existing structures (for example, store derived
3
# arrays in clip.egomotion or add keys to dictionaries), or maintain your own
4
# data alongside and write it in a custom writer stage.

1
from nemo_curator.pipeline import Pipeline
2
from nemo_curator.stages.video.io.video_reader import VideoReader
3
from nemo_curator.stages.video.io.clip_writer import ClipWriterStage
4
5
pipeline = (
6
    Pipeline(name="custom-video-pipeline")
7
    .add_stage(VideoReader(input_video_path="/path/to/videos", video_limit=10))
8
    .add_stage(MyCustomStage())
9
    .add_stage(
10
        ClipWriterStage(
11
            output_path="/path/to/output",
12
            input_path="/path/to/videos",
13
            upload_clips=True,
14
            dry_run=False,
15
            generate_embeddings=False,
16
            generate_previews=False,
17
            generate_captions=False,
18
            embedding_algorithm="cosmos-embed1-224p",
19
            caption_models=["qwen"],
20
            enhanced_caption_models=["qwen_lm"],
21
        )
22
    )
23
)
24
25
# Optionally provide an executor; defaults to XennaExecutor
26
pipeline.run()