> This page is for version 0.5.0 · 26.06.
> For other versions, use one of these documentation indexes:
> - Nightly (default): https://docs.nvidia.com/nemo/automodel/nightly/llms.txt
> - Latest: https://docs.nvidia.com/nemo/automodel/latest/llms.txt
> - 0.5.0 · 26.06: https://docs.nvidia.com/nemo/automodel/v0.5/llms.txt
> - 0.4.0 · 26.04: https://docs.nvidia.com/nemo/automodel/v0.4/llms.txt

> For clean Markdown of any page, append .md to the page URL.
> For a complete documentation index, see https://docs.nvidia.com/nemo/automodel/llms.txt.
> For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://docs.nvidia.com/nemo/automodel/_mcp/server.

# nemo_automodel.components.datasets.diffusion.text_to_image_dataset

## Module Contents

### Classes

| Name                                                                                                                       | Description                                                  |
| -------------------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------ |
| [`TextToImageDataset`](#nemo_automodel-components-datasets-diffusion-text_to_image_dataset-TextToImageDataset)             | Text-to-Image dataset with hierarchical bucket organization. |
| [`TextToImageDatasetConfig`](#nemo_automodel-components-datasets-diffusion-text_to_image_dataset-TextToImageDatasetConfig) | Construction-time configuration for `TextToImageDataset`.    |

### Data

[`PROMPT_IDS_KEY`](#nemo_automodel-components-datasets-diffusion-text_to_image_dataset-PROMPT_IDS_KEY)

[`PROMPT_SUFFIX_IDS_KEY`](#nemo_automodel-components-datasets-diffusion-text_to_image_dataset-PROMPT_SUFFIX_IDS_KEY)

[`PROMPT_TOKEN_ID_KEYS`](#nemo_automodel-components-datasets-diffusion-text_to_image_dataset-PROMPT_TOKEN_ID_KEYS)

[`UNCOND_PROMPT_IDS_KEY`](#nemo_automodel-components-datasets-diffusion-text_to_image_dataset-UNCOND_PROMPT_IDS_KEY)

### API

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

**Bases:** [BaseMultiresolutionDataset](/nemo-automodel/nemo_automodel/components/datasets/diffusion/base_dataset#nemo_automodel-components-datasets-diffusion-base_dataset-BaseMultiresolutionDataset)

Text-to-Image dataset with hierarchical bucket organization.

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

Load a single sample.

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

---

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

Build a `TextToImageDataset` from this `TextToImageDatasetConfig`.

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

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

```python
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).

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