nemo_automodel.components.datasets.diffusion.meta_files_dataset

Module Contents

Classes

Name	Description
`MetaFilesDataset`	PyTorch dataset for WAN2.1 `.meta` files.

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.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

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__)