> 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.image_edit_dataset

Cached, model-agnostic image-edit datasets and dataloaders.

## Module Contents

### Classes

| Name                                                                                                                                    | Description                                                         |
| --------------------------------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------- |
| [`ImageEditDataloaderConfig`](#nemo_automodel-components-datasets-diffusion-image_edit_dataset-ImageEditDataloaderConfig)               | Construction-time configuration for a cached image-edit dataloader. |
| [`ImageEditDataset`](#nemo_automodel-components-datasets-diffusion-image_edit_dataset-ImageEditDataset)                                 | Dataset for cached instruction-based image editing.                 |
| [`ImageEditDatasetConfig`](#nemo_automodel-components-datasets-diffusion-image_edit_dataset-ImageEditDatasetConfig)                     | Construction-time configuration for :class:`ImageEditDataset`.      |
| [`_CompoundBucketSignature`](#nemo_automodel-components-datasets-diffusion-image_edit_dataset-_CompoundBucketSignature)                 | -                                                                   |
| [`_CompoundBucketSignatureMetadata`](#nemo_automodel-components-datasets-diffusion-image_edit_dataset-_CompoundBucketSignatureMetadata) | -                                                                   |
| [`_ImageEditBatch`](#nemo_automodel-components-datasets-diffusion-image_edit_dataset-_ImageEditBatch)                                   | -                                                                   |
| [`_ImageEditBatchMetadata`](#nemo_automodel-components-datasets-diffusion-image_edit_dataset-_ImageEditBatchMetadata)                   | -                                                                   |
| [`_ImageEditCacheMetadata`](#nemo_automodel-components-datasets-diffusion-image_edit_dataset-_ImageEditCacheMetadata)                   | -                                                                   |
| [`_ImageEditSample`](#nemo_automodel-components-datasets-diffusion-image_edit_dataset-_ImageEditSample)                                 | -                                                                   |

### Functions

| Name                                                                                                                                    | Description                                                        |
| --------------------------------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------ |
| [`_collate_image_edit`](#nemo_automodel-components-datasets-diffusion-image_edit_dataset-_collate_image_edit)                           | Collate samples with compatible target and ordered-context shapes. |
| [`_nonnegative_int`](#nemo_automodel-components-datasets-diffusion-image_edit_dataset-_nonnegative_int)                                 | -                                                                  |
| [`_normalize_shape`](#nemo_automodel-components-datasets-diffusion-image_edit_dataset-_normalize_shape)                                 | -                                                                  |
| [`_pad_prompts`](#nemo_automodel-components-datasets-diffusion-image_edit_dataset-_pad_prompts)                                         | Pad cached prompt encodings to a common sequence length.           |
| [`_parse_compound_bucket_signature`](#nemo_automodel-components-datasets-diffusion-image_edit_dataset-_parse_compound_bucket_signature) | -                                                                  |
| [`_require_tensor`](#nemo_automodel-components-datasets-diffusion-image_edit_dataset-_require_tensor)                                   | Read a tensor field from a cache payload.                          |
| [`_signature_to_metadata`](#nemo_automodel-components-datasets-diffusion-image_edit_dataset-_signature_to_metadata)                     | -                                                                  |
| [`_stack_conditioning_tensors`](#nemo_automodel-components-datasets-diffusion-image_edit_dataset-_stack_conditioning_tensors)           | Stack explicitly named model-conditioning tensors.                 |
| [`_validate_latent`](#nemo_automodel-components-datasets-diffusion-image_edit_dataset-_validate_latent)                                 | Validate a cached image latent.                                    |

### Data

[`__all__`](#nemo_automodel-components-datasets-diffusion-image_edit_dataset-__all__)

[`logger`](#nemo_automodel-components-datasets-diffusion-image_edit_dataset-logger)

### API

```python
class nemo_automodel.components.datasets.diffusion.image_edit_dataset.ImageEditDataloaderConfig(
    cache_dir: str,
    quantization: int = 64,
    base_resolution: tuple[int, int] = (1024, 1024),
    drop_last: bool = True,
    shuffle: bool = True,
    dynamic_batch_size: bool = False,
    num_workers: int = 4,
    pin_memory: bool = True,
    prefetch_factor: int = 2,
    seed: int = 42
)
```

Dataclass

Construction-time configuration for a cached image-edit dataloader.

```python
nemo_automodel.components.datasets.diffusion.image_edit_dataset.ImageEditDataloaderConfig.build(
    dp_rank: int,
    dp_world_size: int,
    batch_size: int
) -> nemo_automodel.components.datasets.diffusion.loader.DiffusionDataloaderBuild
```

Build the configured dataset, bucket sampler, and stateful dataloader.

**Parameters:**

Data-parallel rank that consumes the dataloader.

Number of data-parallel ranks.

Base per-rank batch size.

**Returns:** `DiffusionDataloaderBuild`

Materialized stateful dataloader and its resumable bucket sampler.

```python
class nemo_automodel.components.datasets.diffusion.image_edit_dataset.ImageEditDataset(
    cache_dir: str,
    quantization: int = 64
)
```

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

Dataset for cached instruction-based image editing.

Each cache payload contains a target latent `[channels, height, width]`,
an ordered list of context latents with the same axis order, prompt
embeddings `[sequence, hidden]`, an attention mask `[sequence]`, and
an optional explicitly named mapping of model-conditioning tensors.

```python
nemo_automodel.components.datasets.diffusion.image_edit_dataset.ImageEditDataset.__getitem__(
    idx: int
) -> dict[str, object]
```

Load and validate one cached image-edit sample.

**Parameters:**

Zero-based sample index.

**Returns:** `dict[str, object]`

Mapping containing a target latent tensor of shape \[channels,

```python
nemo_automodel.components.datasets.diffusion.image_edit_dataset.ImageEditDataset._group_by_bucket() -> None
```

```python
nemo_automodel.components.datasets.diffusion.image_edit_dataset.ImageEditDataset._load_metadata() -> list[dict[str, object]]
```

```python
nemo_automodel.components.datasets.diffusion.image_edit_dataset.ImageEditDataset._normalize_cache_metadata(
    raw_metadata: object,
    signature: nemo_automodel.components.datasets.diffusion.image_edit_dataset._CompoundBucketSignature,
    prompt_sequence_length: int,
    conditioning_tensors: collections.abc.Mapping[str, torch.Tensor],
    cache_file: pathlib.Path
) -> nemo_automodel.components.datasets.diffusion.image_edit_dataset._ImageEditCacheMetadata
```

Validate and normalize provenance and tensor-shape metadata.

**Parameters:**

Serialized sample provenance and cache metadata.

Target shape \[channels, height, width] and ordered
context shapes \[channels, context\_height, context\_width].

Sequence axis length of prompt embeddings.

Named tensors with arbitrary, explicitly
metadata-declared shapes.

Cache file used to identify invalid fields in errors.

**Returns:** `_ImageEditCacheMetadata`

Validated metadata with spatial shapes, token lengths, and compound

```python
nemo_automodel.components.datasets.diffusion.image_edit_dataset.ImageEditDataset._resolve_contained_path(
    path_value: object,
    field_name: str
) -> pathlib.Path
```

```python
nemo_automodel.components.datasets.diffusion.image_edit_dataset.ImageEditDataset.get_bucket_info() -> dict[str, object]
```

Return target/context shape groups used by the bucket sampler.

**Returns:** `dict[str, object]`

Mapping containing the bucket count and sample count per compound

```python
class nemo_automodel.components.datasets.diffusion.image_edit_dataset.ImageEditDatasetConfig(
    cache_dir: str,
    quantization: int = 64
)
```

Dataclass

Construction-time configuration for :class:`ImageEditDataset`.

Directory containing a preprocessed image-edit cache.

Spatial quantization used for dynamic batch-size calculation.

```python
nemo_automodel.components.datasets.diffusion.image_edit_dataset.ImageEditDatasetConfig.build() -> nemo_automodel.components.datasets.diffusion.image_edit_dataset.ImageEditDataset
```

Build the configured image-edit dataset.

**Returns:** `ImageEditDataset`

Dataset backed by the configured preprocessed cache.

```python
class nemo_automodel.components.datasets.diffusion.image_edit_dataset._CompoundBucketSignature(
    target: tuple[int, int, int],
    contexts: tuple[tuple[int, int, int], ...]
)
```

Dataclass

```python
class nemo_automodel.components.datasets.diffusion.image_edit_dataset._CompoundBucketSignatureMetadata
```

**Bases:** `typing.TypedDict`

```python
class nemo_automodel.components.datasets.diffusion.image_edit_dataset._ImageEditBatch
```

**Bases:** `typing.TypedDict`

```python
class nemo_automodel.components.datasets.diffusion.image_edit_dataset._ImageEditBatchMetadata
```

**Bases:** `typing.TypedDict`

```python
class nemo_automodel.components.datasets.diffusion.image_edit_dataset._ImageEditCacheMetadata
```

**Bases:** `typing.TypedDict`

```python
class nemo_automodel.components.datasets.diffusion.image_edit_dataset._ImageEditSample
```

**Bases:** `typing.TypedDict`

```python
nemo_automodel.components.datasets.diffusion.image_edit_dataset._collate_image_edit(
    batch: list[nemo_automodel.components.datasets.diffusion.image_edit_dataset._ImageEditSample]
) -> nemo_automodel.components.datasets.diffusion.image_edit_dataset._ImageEditBatch
```

Collate samples with compatible target and ordered-context shapes.

**Parameters:**

Samples containing target tensors of shape \[channels, height,
width], ordered context tensors of shape \[channels, context\_height,
context\_width], prompt tensors of shape \[sequence, hidden], prompt
masks of shape \[sequence], and optional named conditioning tensors.

**Returns:** `_ImageEditBatch`

Batch containing `image_latents` of shape \[batch, channels, height,

```python
nemo_automodel.components.datasets.diffusion.image_edit_dataset._nonnegative_int(
    value: object,
    field_name: str
) -> int
```

```python
nemo_automodel.components.datasets.diffusion.image_edit_dataset._normalize_shape(
    value: object,
    field_name: str,
    rank: int,
    allow_zero: bool = False
) -> tuple[int, ...]
```

```python
nemo_automodel.components.datasets.diffusion.image_edit_dataset._pad_prompts(
    embeddings: collections.abc.Sequence[torch.Tensor],
    masks: collections.abc.Sequence[torch.Tensor]
) -> tuple[torch.Tensor, torch.Tensor]
```

Pad cached prompt encodings to a common sequence length.

**Parameters:**

Prompt tensors, each of shape \[sequence, hidden].

Attention-mask tensors, each of shape \[sequence].

**Returns:** `torch.Tensor`

Prompt embeddings of shape \[batch, max\_sequence, hidden] and attention

```python
nemo_automodel.components.datasets.diffusion.image_edit_dataset._parse_compound_bucket_signature(
    value: object,
    field_name: str
) -> nemo_automodel.components.datasets.diffusion.image_edit_dataset._CompoundBucketSignature
```

```python
nemo_automodel.components.datasets.diffusion.image_edit_dataset._require_tensor(
    data: collections.abc.Mapping[str, object],
    field_name: str
) -> torch.Tensor
```

Read a tensor field from a cache payload.

**Parameters:**

Cache mapping whose tensor fields may have arbitrary shapes.

Name of the tensor field to read.

**Returns:** `torch.Tensor`

Cached tensor with arbitrary shape; the field owner validates its layout.

```python
nemo_automodel.components.datasets.diffusion.image_edit_dataset._signature_to_metadata(
    signature: nemo_automodel.components.datasets.diffusion.image_edit_dataset._CompoundBucketSignature
) -> nemo_automodel.components.datasets.diffusion.image_edit_dataset._CompoundBucketSignatureMetadata
```

```python
nemo_automodel.components.datasets.diffusion.image_edit_dataset._stack_conditioning_tensors(
    batch: collections.abc.Sequence[nemo_automodel.components.datasets.diffusion.image_edit_dataset._ImageEditSample]
) -> dict[str, torch.Tensor]
```

Stack explicitly named model-conditioning tensors.

**Parameters:**

Samples whose `conditioning_tensors` entries contain tensors
with metadata-declared arbitrary shapes. A given name must have the
same shape and dtype in every sample.

**Returns:** `dict[str, torch.Tensor]`

Mapping whose tensors have shape \[batch, ...], where trailing axes match

```python
nemo_automodel.components.datasets.diffusion.image_edit_dataset._validate_latent(
    tensor: torch.Tensor,
    field_name: str
) -> tuple[int, int, int]
```

Validate a cached image latent.

**Parameters:**

Tensor of shape \[channels, height, width].

Cache field name used in validation errors.

**Returns:** `tuple[int, int, int]`

Validated tensor shape in \[channels, height, width] order.

```python
nemo_automodel.components.datasets.diffusion.image_edit_dataset.__all__ = ['ImageEditDataloaderConfig', 'ImageEditDataset', 'ImageEditDatasetConfig']
```

```python
nemo_automodel.components.datasets.diffusion.image_edit_dataset.logger = logging.getLogger(__name__)
```