nemo_automodel.components.datasets.diffusion.text_to_image_dataset

Module Contents

Classes

Name	Description
`TextToImageDataset`	Text-to-Image dataset with hierarchical bucket organization.
`TextToImageDatasetConfig`	Construction-time configuration for :class:`TextToImageDataset`.

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.

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

Dataclass

Construction-time configuration for :class:TextToImageDataset.

cache_dir

str

Directory containing preprocessed cache.

train_text_encoder

bool = False

If True, returns tokens instead of embeddings.

nemo_automodel.components.datasets.diffusion.text_to_image_dataset.TextToImageDatasetConfig.build() -> 'TextToImageDataset'

Build a :class:TextToImageDataset from this :class:TextToImageDatasetConfig.