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

Data

PROMPT_IDS_KEY

PROMPT_SUFFIX_IDS_KEY

PROMPT_TOKEN_ID_KEYS

UNCOND_PROMPT_IDS_KEY

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 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 TextToImageDataset from this TextToImageDatasetConfig.

nemo_automodel.components.datasets.diffusion.text_to_image_dataset.PROMPT_IDS_KEY = 'prompt_input_ids'

nemo_automodel.components.datasets.diffusion.text_to_image_dataset.PROMPT_SUFFIX_IDS_KEY = 'prompt_suffix_ids'

nemo_automodel.components.datasets.diffusion.text_to_image_dataset.PROMPT_TOKEN_ID_KEYS = (PROMPT_IDS_KEY, UNCOND_PROMPT_IDS_KEY, PROMPT_SUFFIX_IDS_KEY)

Optional per-sample 1D token-id tensors of models that embed the prompt themselves (no text encoder).

nemo_automodel.components.datasets.diffusion.text_to_image_dataset.UNCOND_PROMPT_IDS_KEY = 'uncond_prompt_input_ids'