nemo_automodel.components.datasets.diffusion.image_edit_dataset

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

Module Contents

Classes

Name	Description
`ImageEditDataloaderConfig`	Construction-time configuration for a cached image-edit dataloader.
`ImageEditDataset`	Dataset for cached instruction-based image editing.
`ImageEditDatasetConfig`	Construction-time configuration for :class:`ImageEditDataset`.
`_CompoundBucketSignature`	-
`_CompoundBucketSignatureMetadata`	-
`_ImageEditBatch`	-
`_ImageEditBatchMetadata`	-
`_ImageEditCacheMetadata`	-
`_ImageEditSample`	-

Functions

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

Data

__all__

logger

API

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.

base_resolution

tuple[int, int] = (1024, 1024)

cache_dir

str

drop_last

bool = True

dynamic_batch_size

bool = False

num_workers

int = 4

pin_memory

bool = True

prefetch_factor

int = 2

quantization

int = 64

seed

int = 42

shuffle

bool = True

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:

dp_rank

int

Data-parallel rank that consumes the dataloader.

dp_world_size

int

Number of data-parallel ranks.

batch_size

int

Base per-rank batch size.

Returns: DiffusionDataloaderBuild

Materialized stateful dataloader and its resumable bucket sampler.

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

Bases: 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.

cache_manifest

dict[str, object] = {}

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

Load and validate one cached image-edit sample.

Parameters:

idx

int

Zero-based sample index.

Returns: dict[str, object]

Mapping containing a target latent tensor of shape [channels,

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

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

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:

raw_metadata

object

Serialized sample provenance and cache metadata.

signature

_CompoundBucketSignature

Target shape [channels, height, width] and ordered context shapes [channels, context_height, context_width].

prompt_sequence_length

int

Sequence axis length of prompt embeddings.

conditioning_tensors

Mapping[str, torch.Tensor]

Named tensors with arbitrary, explicitly metadata-declared shapes.

cache_file

Path

Cache file used to identify invalid fields in errors.

Returns: _ImageEditCacheMetadata

Validated metadata with spatial shapes, token lengths, and compound

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

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

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

Dataclass

Construction-time configuration for :class:ImageEditDataset.

cache_dir

str

Directory containing a preprocessed image-edit cache.

quantization

int = 64

Spatial quantization used for dynamic batch-size calculation.

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.

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

Dataclass

contexts

tuple[tuple[int, int, int], ...]

target

tuple[int, int, int]

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

Bases: typing.TypedDict

contexts

list[list[int]]

target

list[int]

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

Bases: typing.TypedDict

conditioning_tensors

dict[str, Tensor]

context_latents

list[Tensor]

data_type

str

image_latents

Tensor

metadata

_ImageEditBatchMetadata

text_attention_mask

Tensor

text_embeddings

Tensor

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

Bases: typing.TypedDict

batch_size

int

context_token_count

int

context_token_counts

list[int]

samples

list[_ImageEditCacheMetadata]

target_token_count

int

target_token_counts

list[int]

text_token_count

int

text_token_counts

list[int]

total_token_count

int

total_token_counts

list[int]

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

Bases: typing.TypedDict

compound_bucket_signature

_CompoundBucketSignatureMetadata

conditioning_shapes

dict[str, list[int]]

context_spatial_shapes

list[list[int]]

context_token_lengths

list[int]

original_ids

dict[str, str | int]

target_spatial_shape

list[int]

target_token_length

int

text_token_length

int

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

Bases: typing.TypedDict

conditioning_tensors

dict[str, Tensor]

context_latents

list[Tensor]

metadata

_ImageEditCacheMetadata

prompt_attention_mask

Tensor

prompt_embeddings

Tensor

target_latent

Tensor

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:

batch

list[_ImageEditSample]

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,

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

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

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:

embeddings

Sequence[torch.Tensor]

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

masks

Sequence[torch.Tensor]

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

Returns: torch.Tensor

Prompt embeddings of shape [batch, max_sequence, hidden] and attention

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

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:

data

Mapping[str, object]

Cache mapping whose tensor fields may have arbitrary shapes.

field_name

str

Name of the tensor field to read.

Returns: torch.Tensor

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

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

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:

batch

Sequence[_ImageEditSample]

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

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

torch.Tensor

Tensor of shape [channels, height, width].

field_name

str

Cache field name used in validation errors.

Returns: tuple[int, int, int]

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

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

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