Skip to navigation
ReferenceFull Library ReferenceNemo AutomodelNemo AutomodelComponentsDatasetsDiffusionnemo_automodel.components.datasets.diffusion.text_to_video_dataset

nemo_automodel.components.datasets.diffusion.text_to_video_dataset

View as Markdown

Module Contents

Classes

NameDescription
TextToVideoDatasetText-to-Video dataset with multiresolution bucket organization.
TextToVideoDatasetConfigConstruction-time configuration for TextToVideoDataset.

Functions

NameDescription
collate_optional_video_fieldsConcatenate optional video fields present in batch into result dict.
load_optional_video_fieldsExtract optional model-specific fields, moving to device.

Data

VIDEO_OPTIONAL_FIELDS

API

class nemo_automodel.components.datasets.diffusion.text_to_video_dataset.TextToVideoDataset(
    cache_dir: str,
    model_type: str = 'wan',
    device: str = 'cpu'
)

Bases: BaseMultiresolutionDataset

Text-to-Video dataset with multiresolution bucket organization.

Loads preprocessed .meta files organized by resolution bucket. Compatible with SequentialBucketSampler for multiresolution training.

nemo_automodel.components.datasets.diffusion.text_to_video_dataset.TextToVideoDataset.__getitem__(
    idx: int
) -> typing.Dict[str, torch.Tensor]

Load a single video sample from its .meta file.

class nemo_automodel.components.datasets.diffusion.text_to_video_dataset.TextToVideoDatasetConfig(
    cache_dir: str,
    model_type: str = 'wan',
    device: str = 'cpu'
)
Dataclass

Construction-time configuration for TextToVideoDataset.

cache_dir
str

Directory containing preprocessed cache (metadata.json + shards + WxH/*.meta).

device
str = 'cpu'

Device to load tensors to.

model_type
str = 'wan'

Model type for model-specific fields (e.g. ‘wan’, ‘hunyuan’).

nemo_automodel.components.datasets.diffusion.text_to_video_dataset.TextToVideoDatasetConfig.build() -> 'TextToVideoDataset'

Build a TextToVideoDataset from this TextToVideoDatasetConfig.

nemo_automodel.components.datasets.diffusion.text_to_video_dataset.collate_optional_video_fields(
    batch: typing.List[typing.Dict],
    result: dict
) -> None

Concatenate optional video fields present in batch into result dict.

nemo_automodel.components.datasets.diffusion.text_to_video_dataset.load_optional_video_fields(
    data: dict,
    device: str = 'cpu'
) -> dict

Extract optional model-specific fields, moving to device.

nemo_automodel.components.datasets.diffusion.text_to_video_dataset.VIDEO_OPTIONAL_FIELDS = ('text_mask', 'text_embeddings_2', 'text_mask_2', 'image_embeds', 'audio_latents...