nemo_automodel.components.datasets.diffusion.meta_files_dataset

Module Contents

Classes

Name	Description
`MetaFilesDataloaderConfig`	Construction-time configuration for a pre-encoded diffusion dataloader.
`MetaFilesDataset`	PyTorch dataset for WAN2.1 `.meta` files.
`MetaFilesDatasetConfig`	Construction-time configuration for :class:`MetaFilesDataset`.

Functions

Name	Description
`build_dataloader`	Build a dataloader for pre-encoded diffusion metadata files.
`build_node_parallel_sampler`	Build a distributed sampler when torch.distributed is initialized.
`collate_fn`	Collate encoded video metadata samples into a training batch.
`create_dataloader`	Create a default metadata dataloader for node-parallel loading.

Data

logger

API

class nemo_automodel.components.datasets.diffusion.meta_files_dataset.MetaFilesDataloaderConfig(
    meta_folder: str,
    shuffle: bool = True,
    num_workers: int = 2,
    device: str = 'cpu',
    transform_text: typing.Callable[[torch.Tensor], torch.Tensor] | None = None,
    transform_video: typing.Callable[[torch.Tensor], torch.Tensor] | None = None,
    filter_fn: typing.Callable[[dict[str, object]], bool] | None = None,
    max_files: int | None = None
)

Dataclass

Construction-time configuration for a pre-encoded diffusion dataloader.

device

str = 'cpu'

filter_fn

Callable[[dict[str, object]], bool] | None = None

max_files

int | None = None

meta_folder

str

num_workers

int = 2

shuffle

bool = True

transform_text

Callable[[torch.Tensor], Tensor] | None = None

transform_video

Callable[[torch.Tensor], Tensor] | None = None

nemo_automodel.components.datasets.diffusion.meta_files_dataset.MetaFilesDataloaderConfig.build(
    dp_rank: int,
    dp_world_size: int,
    batch_size: int
) -> nemo_automodel.components.datasets.diffusion.loader.DiffusionDataloaderBuild

Build the configured metadata dataset, sampler, and dataloader.

class nemo_automodel.components.datasets.diffusion.meta_files_dataset.MetaFilesDataset(
    meta_folder: str,
    transform_text: typing.Optional[typing.Callable[[torch.Tensor], torch.Tensor]] = None,
    transform_video: typing.Optional[typing.Callable[[torch.Tensor], torch.Tensor]] = None,
    filter_fn: typing.Optional[typing.Callable[[Dict], bool]] = None,
    device: str = 'cpu',
    max_files: typing.Optional[int] = None
)

Bases: Dataset

PyTorch dataset for WAN2.1 .meta files.

meta_files

= sorted(self.meta_folder.glob('*.meta'))

meta_folder

= Path(meta_folder)

nemo_automodel.components.datasets.diffusion.meta_files_dataset.MetaFilesDataset.__getitem__(
    index: int
) -> typing.Dict[str, torch.Tensor]

nemo_automodel.components.datasets.diffusion.meta_files_dataset.MetaFilesDataset.__len__() -> int

nemo_automodel.components.datasets.diffusion.meta_files_dataset.MetaFilesDataset._log_dataset_stats() -> None

class nemo_automodel.components.datasets.diffusion.meta_files_dataset.MetaFilesDatasetConfig(
    meta_folder: str,
    device: str = 'cpu',
    max_files: int | None = None
)

Dataclass

Construction-time configuration for :class:MetaFilesDataset.

device

str = 'cpu'

Device to load tensors to.

max_files

int | None = None

Maximum number of .meta files to use (None means no limit).

meta_folder

str

Path to the folder containing .meta files.

nemo_automodel.components.datasets.diffusion.meta_files_dataset.MetaFilesDatasetConfig.build(
    transform_text: typing.Callable[[torch.Tensor], torch.Tensor] | None = None,
    transform_video: typing.Callable[[torch.Tensor], torch.Tensor] | None = None,
    filter_fn: typing.Callable[[dict[str, object]], bool] | None = None
) -> 'MetaFilesDataset'

Build a :class:MetaFilesDataset from this :class:MetaFilesDatasetConfig.

Parameters:

transform_text

Callable[[torch.Tensor], torch.Tensor] | NoneDefaults to None

Optional transform of text embeddings shaped [1, S, E], where S is text sequence length and E is embedding width. It must return the same semantic axis order.

transform_video

Callable[[torch.Tensor], torch.Tensor] | NoneDefaults to None

Optional transform of video latents shaped [1, C, T, Y, X], where C is latent channels, T is latent frames, and Y/X are spatial height/width. It must return the same semantic axis order.

filter_fn

Callable[[dict[str, object]], bool] | NoneDefaults to None

Optional callable that filters samples by their metadata dict.

Returns: 'MetaFilesDataset'

Dataset loading tensors on the configured device.

nemo_automodel.components.datasets.diffusion.meta_files_dataset.build_dataloader(
    meta_folder: str,
    batch_size: int,
    dp_rank: int,
    dp_world_size: int,
    shuffle: bool = True,
    num_workers: int = 2,
    device: str = 'cpu',
    transform_text: typing.Optional[typing.Callable[[torch.Tensor], torch.Tensor]] = None,
    transform_video: typing.Optional[typing.Callable[[torch.Tensor], torch.Tensor]] = None,
    filter_fn: typing.Optional[typing.Callable[[Dict], bool]] = None,
    max_files: typing.Optional[int] = None
) -> typing.Tuple[torch.utils.data.DataLoader, typing.Optional[torch.utils.data.DistributedSampler]]

Build a dataloader for pre-encoded diffusion metadata files.

nemo_automodel.components.datasets.diffusion.meta_files_dataset.build_node_parallel_sampler(
    dataset: 'Dataset',
    dp_rank: int,
    dp_world_size: int,
    shuffle: bool = True
) -> typing.Optional['DistributedSampler']

Build a distributed sampler when torch.distributed is initialized.

nemo_automodel.components.datasets.diffusion.meta_files_dataset.collate_fn(
    batch: typing.List[typing.Dict[str, torch.Tensor]]
) -> typing.Dict[str, torch.Tensor]

Collate encoded video metadata samples into a training batch.

nemo_automodel.components.datasets.diffusion.meta_files_dataset.create_dataloader(
    meta_folder: str,
    batch_size: int,
    num_nodes: int
) -> typing.Tuple[torch.utils.data.DataLoader, typing.Optional[torch.utils.data.DistributedSampler]]

Create a default metadata dataloader for node-parallel loading.

nemo_automodel.components.datasets.diffusion.meta_files_dataset.logger = logging.getLogger(__name__)