nemo_automodel.components.datasets.diffusion.base_dataset#

Module Contents#

Classes#

BaseMultiresolutionDataset

Abstract base class for multiresolution datasets with bucket-based sampling.

Data#

logger

API#

nemo_automodel.components.datasets.diffusion.base_dataset.logger#

‘getLogger(…)’

class nemo_automodel.components.datasets.diffusion.base_dataset.BaseMultiresolutionDataset(cache_dir: str, quantization: int = 64)#

Bases: torch.utils.data.Dataset, abc.ABC

Abstract base class for multiresolution datasets with bucket-based sampling.

Initialization

Parameters:

  • cache_dir – Directory containing preprocessed cache (metadata.json + shards)

  • quantization – Resolution quantization factor (64 for images, 8 for video)

_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#
abstractmethod __getitem__(idx: int) → Dict#

Load a single sample. Subclasses must implement.