nemo_automodel.components.datasets.diffusion.text_to_image_dataset#

Module Contents#

Classes#

TextToImageDataset

Text-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: nemo_automodel.components.datasets.diffusion.base_dataset.BaseMultiresolutionDataset

Text-to-Image dataset with hierarchical bucket organization.

Initialization

Parameters:

  • cache_dir – Directory containing preprocessed cache

  • train_text_encoder – If True, returns tokens instead of embeddings

__getitem__(idx: int) → Dict[str, torch.Tensor]#

Load a single sample.