Method __init__

utils/dataset.py:207–231

(self, metadata_dataset, directory_config, size_bucket, cache_base, directory_dataset)

# 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}')