DevelopmentAPI ReferenceFull Library ReferenceNemo AutomodelNemo AutomodelComponentsDatasetsDiffusion

nemo_automodel.components.datasets.diffusion.text_to_image_dataset

View as Markdown

Module Contents

Classes

NameDescription
TextToImageDatasetText-to-Image dataset with hierarchical bucket organization.

API

class nemo_automodel.components.datasets.diffusion.text_to_image_dataset.TextToImageDataset(
    cache_dir: str,
    train_text_encoder: bool = False
)

Bases: BaseMultiresolutionDataset

Text-to-Image dataset with hierarchical bucket organization.

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

Load a single sample.