nemo_automodel.components.datasets.diffusion.text_to_image_dataset#

Module Contents#

Classes#

TextToImageDataset

Text-to-Image dataset with hierarchical bucket organization.

Data#

logger

API#

nemo_automodel.components.datasets.diffusion.text_to_image_dataset.logger#

‘getLogger(…)’

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

Bases: torch.utils.data.Dataset

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

_load_metadata() → List[Dict]#

Load metadata from cache directory.

Expects metadata.json with “shards” key referencing shard files.

_aspect_ratio_to_name(aspect_ratio: float) → str#

Convert aspect ratio to a descriptive name.

_group_by_bucket()#

Group samples by bucket (aspect_ratio + resolution).

get_bucket_info() → Dict#

Get bucket organization information.

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

Load a single sample.