Class SizeBucketDataset

utils/dataset.py:206–336

# The smallest unit of a dataset. Represents a single size bucket from a single folder of images
# and captions on disk. Not batched; returns individual items.
class SizeBucketDataset:
    def __init__(self, metadata_dataset, directory_config, size_bucket, cache_base, directory_dataset):
        # Shuffle deterministically based on size bucket, so that two resolutions of the same aspect ratio get different
        # orders, which mixes data better when training on multiple resolutions at once.
        seed = seed_from_hash(size_bucket)
        metadata_dataset = metadata_dataset.shuffle(seed=seed)
        self.metadata_dataset = metadata_dataset
        self.directory_config = directory_config
        self.size_bucket = size_bucket
        self.path = Path(self.directory_config['path'])
        self.cache_dir = cache_base / f'cache_{bucket_suffix(size_bucket)}'
        self.captions_dict = directory_dataset.captions_dict  # optional

        if len(size_bucket) == 4:
            # rename old folder name to the new one for convenience
            old_cache_dir = cache_base / f'cache_{bucket_suffix(size_bucket[1:])}'
            if old_cache_dir.exists() and not self.cache_dir.exists():
                old_cache_dir.rename(self.cache_dir)

        os.makedirs(self.cache_dir, exist_ok=True)
        self.text_embedding_datasets = []
        self.uncond_text_embeddings = []
        self.num_repeats = self.directory_config['num_repeats']
        self.shuffle_skip = max(directory_config.get('cache_shuffle_num', 0), 1) # Should be provided in DirectoryDataset
        if self.num_repeats <= 0:
            raise ValueError(f'num_repeats must be >0, was {self.num_repeats}')

    def cache_latents(self, map_fn, regenerate_cache=False, trust_cache=False, caching_batch_size=1):
        print(f'caching latents: {self.size_bucket}')
        self.latent_dataset = _map_and_cache(
            self.metadata_dataset,
            map_fn,
            self.cache_dir,
            cache_file_prefix='latents_',
            regenerate_cache=regenerate_cache,
            caching_batch_size=caching_batch_size,
        )
        assert len(self.latent_dataset) == len(self.metadata_dataset), (len(self.latent_dataset), len(self.metadata_dataset))

        iteration_order_cache_dir = self.cache_dir / 'iteration_order'

        if regenerate_cache or not iteration_order_cache_dir.exists() or not trust_cache:
            print('Building iteration order')
            image_spec_to_latents_idx = {
                tuple(image_spec): i
                for i, image_spec in enumerate(self.metadata_dataset['image_spec'])
            }

            equal_num_captions = True
            num_captions = None
            for example in self.metadata_dataset.select_columns(['caption']):
                n = len(example['caption'])
                if num_captions is not None and n != num_captions:
                    equal_num_captions = False
                    break
                num_captions = n

            if equal_num_captions: