nemo_automodel.components.datasets.diffusion.text_to_video_dataset

Module Contents

Classes

Name	Description
`TextToVideoDataset`	Text-to-Video dataset with multiresolution bucket organization.
`TextToVideoDatasetConfig`	Construction-time configuration for :class:`TextToVideoDataset`.

Functions

Name	Description
`collate_optional_video_fields`	Concatenate optional video fields present in batch into result dict.
`load_optional_video_fields`	Extract 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 :class: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 :class:TextToVideoDataset from this :class: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...