Add ability to merge dataset buckets. Useful when using numerous different datasets.

jobs/process/BaseSDTrainProcess.py

@@ -2244,10 +2244,15 @@ class BaseSDTrainProcess(BaseTrainProcess):
                         ds.cache_tensors_to_disk = True
         # load datasets if passed in the root process
         if self.datasets is not None:
-            self.data_loader = get_dataloader_from_datasets(self.datasets, self.train_config.batch_size, self.sd)
+            self.data_loader = get_dataloader_from_datasets(
+                self.datasets, self.train_config.batch_size, self.sd,
+                pool_datasets=self.train_config.pool_datasets,
+            )
         if self.datasets_reg is not None:
-            self.data_loader_reg = get_dataloader_from_datasets(self.datasets_reg, self.train_config.batch_size,
-                                                                self.sd)
+            self.data_loader_reg = get_dataloader_from_datasets(
+                self.datasets_reg, self.train_config.batch_size, self.sd,
+                pool_datasets=self.train_config.pool_datasets,
+            )
 
         flush()
         self.last_save_step = self.step_num

toolkit/config_modules.py

@@ -568,6 +568,9 @@ class TrainConfig:
         self.unload_text_encoder = kwargs.get('unload_text_encoder', False)
         # will toggle all datasets to cache text embeddings
         self.cache_text_embeddings: bool = kwargs.get('cache_text_embeddings', False)
+        # merge all datasets into one file list so buckets/batches mix across folders;
+        # reg datasets pool separately from non-reg ones
+        self.pool_datasets: bool = kwargs.get('pool_datasets', False)
         # for swapping which parameters are trained during training
         self.do_paramiter_swapping = kwargs.get('do_paramiter_swapping', False)
         # 0.1 is 10% of the parameters active at a time lower is less vram, higher is more
@@ -1109,10 +1112,28 @@ class DatasetConfig:
 
 def preprocess_dataset_raw_config(raw_config: List[dict]) -> List[dict]:
     """
-    This just splits up the datasets by resolutions so you dont have to do it manually
+    This just splits up the datasets by paths and resolutions so you dont have to do it manually
     :param raw_config:
     :return:
     """
+    # split up datasets by folder_path / dataset_path lists: one dataset per path,
+    # all other settings shared (pair with train.pool_datasets to mix them)
+    path_split_config = []
+    for dataset in raw_config:
+        split = False
+        for path_key in ['folder_path', 'dataset_path']:
+            paths = dataset.get(path_key, None)
+            if isinstance(paths, (list, tuple)):
+                for path in paths:
+                    dataset_copy = dataset.copy()
+                    dataset_copy[path_key] = path
+                    path_split_config.append(dataset_copy)
+                split = True
+                break
+        if not split:
+            path_split_config.append(dataset)
+    raw_config = path_split_config
+
     # split up datasets by resolutions
     new_config = []
     for dataset in raw_config:

toolkit/data_loader.py

@@ -390,11 +390,17 @@ class AiToolkitDataset(LatentCachingMixin, ControlCachingMixin, CLIPCachingMixin
             dataset_config: 'DatasetConfig',
             batch_size=1,
             sd: 'StableDiffusion' = None,
+            source_configs: List['DatasetConfig'] = None,
     ):
+        # a pooled dataset scans several configs into ONE file_list so buckets and
+        # batches mix across folders. Single-source is the same code with one config.
+        self.source_configs: List['DatasetConfig'] = list(source_configs) if source_configs else [dataset_config]
+        self.is_pooled = len(self.source_configs) > 1
         self.dataset_config = dataset_config
         # update bucket divisibility
-        self.dataset_config.bucket_tolerance = sd.get_bucket_divisibility()
-        self.is_video = dataset_config.num_frames > 1 or dataset_config.auto_frame_count
+        for config in self.source_configs:
+            config.bucket_tolerance = sd.get_bucket_divisibility()
+        self.is_video = any(c.num_frames > 1 or c.auto_frame_count for c in self.source_configs)
         self.is_audio_model = hasattr(sd, 'is_audio_model') and sd.is_audio_model if sd is not None else False
         # text-generating multimodal models take audio, image and video files in one dataset
         self.is_multimodal_llm = getattr(sd, 'is_multimodal_llm', False) if sd is not None else False
@@ -403,12 +409,16 @@ class AiToolkitDataset(LatentCachingMixin, ControlCachingMixin, CLIPCachingMixin
         self.dataset_path = dataset_config.dataset_path
         if self.dataset_path is None:
             self.dataset_path = folder_path
+        if self.is_pooled:
+            self.dataset_path = f"pool of {len(self.source_configs)} datasets"
 
-        self.is_caching_latents = dataset_config.cache_latents or dataset_config.cache_latents_to_disk
-        self.is_caching_latents_to_memory = dataset_config.cache_latents
-        self.is_caching_latents_to_disk = dataset_config.cache_latents_to_disk
-        self.is_caching_clip_vision_to_disk = dataset_config.cache_clip_vision_to_disk
-        self.is_generating_controls = len(dataset_config.controls) > 0
+        self.is_caching_latents = any(c.cache_latents or c.cache_latents_to_disk for c in self.source_configs)
+        self.is_caching_latents_to_memory = any(c.cache_latents for c in self.source_configs)
+        self.is_caching_latents_to_disk = any(c.cache_latents_to_disk for c in self.source_configs)
+        self.is_caching_clip_vision_to_disk = any(c.cache_clip_vision_to_disk for c in self.source_configs)
+        self.is_generating_controls = any(len(c.controls) > 0 for c in self.source_configs)
+        # mixin init read this off dataset_config only
+        self.is_caching_text_embeddings = any(c.cache_text_embeddings for c in self.source_configs)
         self.epoch_num = 0
 
         self.sd = sd
@@ -428,9 +438,45 @@ class AiToolkitDataset(LatentCachingMixin, ControlCachingMixin, CLIPCachingMixin
         self.resolution = dataset_config.resolution
         self.caption_dict = None
         self.file_list: List['FileItemDTO'] = []
+        self.transform = self._build_transform(dataset_config)
+
+        for config in self.source_configs:
+            self._scan_source(config)
+
+        if self.is_pooled:
+            print_acc(f"Pooled {len(self.source_configs)} datasets into {len(self.file_list)} items")
+
+        self.setup_epoch()
+
+    def _build_transform(self, dataset_config: 'DatasetConfig'):
+        if dataset_config.standardize_images:
+            if self.sd.is_xl or self.sd.is_vega or self.sd.is_ssd:
+                NormalizeMethod = NormalizeSDXLTransform
+            else:
+                NormalizeMethod = NormalizeSD15Transform
+
+            return transforms.Compose([
+                transforms.ToTensor(),
+                RescaleTransform(),
+                NormalizeMethod(),
+            ])
+        else:
+            return transforms.Compose([
+                transforms.ToTensor(),
+                RescaleTransform(),
+            ])
+
+    def _scan_source(self, dataset_config: 'DatasetConfig'):
+        # scans one config (folder or json) into FileItemDTOs and appends them to file_list.
+        # size database, captions and controls stay per source folder.
+        dataset_path = dataset_config.dataset_path
+        if dataset_path is None:
+            dataset_path = dataset_config.folder_path
+        transform = self.transform if dataset_config is self.dataset_config else self._build_transform(dataset_config)
+        caption_dict = None
 
         # check if dataset_path is a folder or json
-        if os.path.isdir(self.dataset_path):
+        if os.path.isdir(dataset_path):
             extensions = image_extensions
             if self.is_audio_model:
                 # only look for audio files
@@ -443,48 +489,34 @@ class AiToolkitDataset(LatentCachingMixin, ControlCachingMixin, CLIPCachingMixin
                 extensions = video_extensions + image_extensions
             # prune hidden dirs (.thumbs, .tmp) so their contents never train
             file_list = []
-            for root, dirs, files in os.walk(self.dataset_path):
+            for root, dirs, files in os.walk(dataset_path):
                 dirs[:] = [d for d in dirs if not d.startswith('.')]
                 file_list.extend(os.path.join(root, file) for file in files if file.lower().endswith(tuple(extensions)) and not file.startswith('.'))
         else:
             # assume json
-            with open(self.dataset_path, 'r') as f:
-                self.caption_dict = json.load(f)
+            with open(dataset_path, 'r') as f:
+                caption_dict = json.load(f)
                 # keys are file paths
-                file_list = list(self.caption_dict.keys())
+                file_list = list(caption_dict.keys())
+            if self.caption_dict is None:
+                self.caption_dict = {}
+            self.caption_dict.update(caption_dict)
 
         # remove items in the _controls_ folder
         file_list = [x for x in file_list if not os.path.basename(os.path.dirname(x)) == "_controls"]
 
-        if self.dataset_config.num_repeats > 1:
+        if dataset_config.num_repeats > 1:
             # repeat the list
-            file_list = file_list * self.dataset_config.num_repeats
-
-        if self.dataset_config.standardize_images:
-            if self.sd.is_xl or self.sd.is_vega or self.sd.is_ssd:
-                NormalizeMethod = NormalizeSDXLTransform
-            else:
-                NormalizeMethod = NormalizeSD15Transform
-
-            self.transform = transforms.Compose([
-                transforms.ToTensor(),
-                RescaleTransform(),
-                NormalizeMethod(),
-            ])
-        else:
-            self.transform = transforms.Compose([
-                transforms.ToTensor(),
-                RescaleTransform(),
-            ])
+            file_list = file_list * dataset_config.num_repeats
 
         # this might take a while
-        print_acc(f"Dataset: {self.dataset_path}")
+        print_acc(f"Dataset: {dataset_path}")
         if self.is_video:
             print_acc(f"  -  Preprocessing video dimensions")
         else:
             print_acc(f"  -  Preprocessing image dimensions")
-        dataset_folder = self.dataset_path
-        if not os.path.isdir(self.dataset_path):
+        dataset_folder = dataset_path
+        if not os.path.isdir(dataset_path):
             dataset_folder = os.path.dirname(dataset_folder)
         
         dataset_size_file = os.path.join(dataset_folder, '.aitk_size.json')
@@ -492,20 +524,21 @@ class AiToolkitDataset(LatentCachingMixin, ControlCachingMixin, CLIPCachingMixin
         if os.path.exists(dataset_size_file):
             try:
                 with open(dataset_size_file, 'r') as f:
-                    self.size_database = json.load(f)
+                    size_database = json.load(f)
                 
-                if "__version__" not in self.size_database or self.size_database["__version__"] != dataloader_version:
+                if "__version__" not in size_database or size_database["__version__"] != dataloader_version:
                     print_acc("Upgrading size database to new version")
                     # old version, delete and recreate
-                    self.size_database = {}
+                    size_database = {}
             except Exception as e:
                 print_acc(f"Error loading size database: {dataset_size_file}")
                 print_acc(e)
-                self.size_database = {}
+                size_database = {}
         else:
-            self.size_database = {}
+            size_database = {}
         
-        self.size_database["__version__"] = dataloader_version
+        size_database["__version__"] = dataloader_version
+        self.size_database = size_database
 
         # cache keys come from the model so a model can invalidate them on its own kwargs
         latent_space_version = self.sd.get_latent_space_version() if self.sd is not None else "sd1"
@@ -518,6 +551,7 @@ class AiToolkitDataset(LatentCachingMixin, ControlCachingMixin, CLIPCachingMixin
             if hasattr(self.sd.unet, 'config') and hasattr(self.sd.unet.config, 'temporal_compression_ratio'):
                 temporal_compression = self.sd.unet.config.temporal_compression_ratio
         
+        source_items: List['FileItemDTO'] = []
         bad_count = 0
         for file in tqdm(file_list):
             try:
@@ -526,8 +560,8 @@ class AiToolkitDataset(LatentCachingMixin, ControlCachingMixin, CLIPCachingMixin
                     path=file,
                     is_audio_model=self.is_audio_model or (self.is_multimodal_llm and file.lower().endswith(tuple(audio_extensions))),
                     dataset_config=dataset_config,
-                    dataloader_transforms=self.transform,
-                    size_database=self.size_database,
+                    dataloader_transforms=transform,
+                    size_database=size_database,
                     dataset_root=dataset_folder,
                     encode_control_in_text_embeddings=self.sd.encode_control_in_text_embeddings if self.sd else False,
                     encode_first_frame_in_text_embeddings=getattr(self.sd, 'encode_first_frame_in_text_embeddings', False) if self.sd else False,
@@ -539,7 +573,7 @@ class AiToolkitDataset(LatentCachingMixin, ControlCachingMixin, CLIPCachingMixin
                     temporal_compression=temporal_compression,
                     sample_rate=self.sd.sample_rate if (self.is_audio_model or self.is_multimodal_llm) and self.sd is not None else 48000,
                 )
-                self.file_list.append(file_item)
+                source_items.append(file_item)
             except Exception as e:
                 print_acc(traceback.format_exc())
                 if self.is_video:
@@ -551,47 +585,47 @@ class AiToolkitDataset(LatentCachingMixin, ControlCachingMixin, CLIPCachingMixin
 
         # save the size database
         with open(dataset_size_file, 'w') as f:
-            json.dump(self.size_database, f)
+            json.dump(size_database, f)
         
         if self.is_video:
-            num_videos = len([x for x in self.file_list if x.is_video])
-            num_images = len(self.file_list) - num_videos
+            num_videos = len([x for x in source_items if x.is_video])
+            num_images = len(source_items) - num_videos
             if num_images > 0:
                 print_acc(f"  -  Found {num_videos} videos and {num_images} images")
             else:
                 print_acc(f"  -  Found {num_videos} videos")
-            assert len(self.file_list) > 0, f"no videos found in {self.dataset_path}"
+            assert len(source_items) > 0, f"no videos found in {dataset_path}"
         else:
-            print_acc(f"  -  Found {len(self.file_list)} images")
-            assert len(self.file_list) > 0, f"no images found in {self.dataset_path}"
+            print_acc(f"  -  Found {len(source_items)} images")
+            assert len(source_items) > 0, f"no images found in {dataset_path}"
 
         # handle x axis flips
-        if self.dataset_config.flip_x:
+        if dataset_config.flip_x:
             print_acc("  -  adding x axis flips")
-            current_file_list = [x for x in self.file_list]
+            current_file_list = [x for x in source_items]
             for file_item in current_file_list:
                 # create a copy that is flipped on the x axis
                 new_file_item = copy.deepcopy(file_item)
                 new_file_item.flip_x = True
-                self.file_list.append(new_file_item)
+                source_items.append(new_file_item)
 
         # handle y axis flips
-        if self.dataset_config.flip_y:
+        if dataset_config.flip_y:
             print_acc("  -  adding y axis flips")
-            current_file_list = [x for x in self.file_list]
+            current_file_list = [x for x in source_items]
             for file_item in current_file_list:
                 # create a copy that is flipped on the y axis
                 new_file_item = copy.deepcopy(file_item)
                 new_file_item.flip_y = True
-                self.file_list.append(new_file_item)
+                source_items.append(new_file_item)
 
-        if self.dataset_config.flip_x or self.dataset_config.flip_y:
+        if dataset_config.flip_x or dataset_config.flip_y:
             if self.is_video:
-                print_acc(f"  -  Found {len(self.file_list)} videos after adding flips")
+                print_acc(f"  -  Found {len(source_items)} videos after adding flips")
             else:
-                print_acc(f"  -  Found {len(self.file_list)} images after adding flips")
+                print_acc(f"  -  Found {len(source_items)} images after adding flips")
 
-        self.setup_epoch()
+        self.file_list.extend(source_items)
 
     def setup_epoch(self):
         if self.epoch_num == 0:
@@ -639,7 +673,7 @@ class AiToolkitDataset(LatentCachingMixin, ControlCachingMixin, CLIPCachingMixin
     def _get_single_item(self, index, _attempts=0) -> 'FileItemDTO':
         file_item: 'FileItemDTO' = copy.deepcopy(self.file_list[index])
         try:
-            file_item.load_and_process_image(self.transform)
+            file_item.load_and_process_image(file_item.dataloader_transforms or self.transform)
         except Exception as e:
             print(f"Error loading image, skipping and loading a different one: {file_item.path} ({e})")
             if _attempts >= 10:
@@ -681,6 +715,7 @@ def get_dataloader_from_datasets(
         dataset_options,
         batch_size=1,
         sd: 'StableDiffusion' = None,
+        pool_datasets: bool = False,
 ) -> DataLoader:
     if dataset_options is None or len(dataset_options) == 0:
         return None
@@ -701,18 +736,34 @@ def get_dataloader_from_datasets(
                 dataset_config_list.append(DatasetConfig(**x))
 
     for config in dataset_config_list:
-
-        if config.type == 'image':
-            # dataset level batch_size overrides the train config batch_size when set
-            dataset_batch_size = config.batch_size if config.batch_size is not None else batch_size
-            dataset = AiToolkitDataset(config, batch_size=dataset_batch_size, sd=sd)
-            datasets.append(dataset)
+        if config.type != 'image':
+            raise ValueError(f"invalid dataset type: {config.type}")
         if config.buckets:
             has_buckets = True
         if config.cache_latents or config.cache_latents_to_disk:
             is_caching_latents = True
+
+    if pool_datasets and len(dataset_config_list) > 1:
+        # one dataset over every config: buckets and batches mix across folders.
+        # callers already split reg and non-reg into separate loaders.
+        batch_sizes = {c.batch_size if c.batch_size is not None else batch_size for c in dataset_config_list}
+        if len(batch_sizes) > 1:
+            raise ValueError(
+                f"pool_datasets requires every dataset to use the same batch_size, got {sorted(batch_sizes)}"
+            )
+        if len({c.buckets for c in dataset_config_list}) > 1:
+            raise ValueError("pool_datasets requires buckets to be enabled or disabled on every dataset")
+        datasets.append(AiToolkitDataset(
+            dataset_config_list[0],
+            batch_size=batch_sizes.pop(),
+            sd=sd,
+            source_configs=dataset_config_list,
+        ))
     else:
-            raise ValueError(f"invalid dataset type: {config.type}")
+        for config in dataset_config_list:
+            # dataset level batch_size overrides the train config batch_size when set
+            dataset_batch_size = config.batch_size if config.batch_size is not None else batch_size
+            datasets.append(AiToolkitDataset(config, batch_size=dataset_batch_size, sd=sd))
 
     concatenated_dataset = ConcatDataset(datasets)
 

toolkit/dataloader_mixins.py

@@ -192,6 +192,43 @@ class Bucket:
         self.file_list_idx: List[int] = []
 
 
+def get_batch_signature(config: 'DatasetConfig') -> tuple:
+    # every dataset-level value the batch DTO / trainer assumes is constant across a batch
+    control_path = config.control_path
+    if control_path is None:
+        num_control_paths = 0
+    elif isinstance(control_path, list):
+        num_control_paths = len(control_path)
+    else:
+        num_control_paths = 1
+    return (
+        num_control_paths,
+        tuple(config.controls),
+        config.control_from_same_folder,
+        config.num_controls_from_same_folder,
+        config.inpaint_path is not None,
+        config.mask_path is not None,
+        config.alpha_mask,
+        config.unconditional_path is not None,
+        config.clip_image_path is not None,
+        config.clip_image_from_same_folder,
+        config.cache_latents,
+        config.cache_latents_to_disk,
+        config.cache_tensors_to_disk,
+        config.cache_text_embeddings,
+        config.load_image_when_caching_latents,
+        config.standardize_images,
+        config.num_frames,
+        config.auto_frame_count,
+        config.fps,
+        config.do_i2v,
+        config.do_audio,
+        config.guidance_type,
+        config.diff_output_preservation,
+        len(config.extra_values),
+    )
+
+
 class BucketsMixin:
     def __init__(self):
         self.buckets: Dict[str, Bucket] = {}
@@ -227,23 +264,33 @@ class BucketsMixin:
         self.buckets = {}  # clear it
 
         config: 'DatasetConfig' = self.dataset_config
-        resolution = config.resolution
         bucket_tolerance = config.bucket_tolerance
         file_list: List['FileItemDTO'] = self.file_list
+        # pooled datasets suffix the key so only batch-compatible sources share a bucket
+        is_pooled = getattr(self, 'is_pooled', False)
+        signature_ids: Dict[tuple, int] = {}
 
         # for file_item in enumerate(file_list):
         for idx, file_item in enumerate(file_list):
             file_item: 'FileItemDTO' = file_item
+            item_config: 'DatasetConfig' = file_item.dataset_config
+            key_suffix = ''
+            if is_pooled:
+                sig = get_batch_signature(item_config)
+                if sig not in signature_ids:
+                    signature_ids[sig] = len(signature_ids)
+                key_suffix = f'#s{signature_ids[sig]}'
             if file_item.is_audio_model:
-                bucket_key = f"{file_item.width}ms"
+                bucket_key = f"{file_item.width}ms{key_suffix}"
                 if bucket_key not in self.buckets:
                     self.buckets[bucket_key] = Bucket(file_item.width, 1)
                 self.buckets[bucket_key].file_list_idx.append(idx)
                 continue
-            width = int(file_item.width * file_item.dataset_config.scale)
-            height = int(file_item.height * file_item.dataset_config.scale)
+            resolution = item_config.resolution
+            width = int(file_item.width * item_config.scale)
+            height = int(file_item.height * item_config.scale)
 
-            if self.dataset_config.square_crop:
+            if item_config.square_crop:
                 # we scale first so smallest size matches resolution
                 scale_factor_x = resolution / width
                 scale_factor_y = resolution / height
@@ -282,7 +329,7 @@ class BucketsMixin:
                 new_width = bucket_resolution["width"]
                 new_height = bucket_resolution["height"]
 
-                if self.dataset_config.random_crop:
+                if item_config.random_crop:
                     # random crop
                     crop_x = random.randint(0, file_item.scale_to_width - new_width)
                     crop_y = random.randint(0, file_item.scale_to_height - new_height)
@@ -301,6 +348,7 @@ class BucketsMixin:
             if self.is_video:
                 # images (1 frame) and videos must not mix in a batch
                 bucket_key += f'x{file_item.num_frames}f'
+            bucket_key += key_suffix
             if bucket_key not in self.buckets:
                 self.buckets[bucket_key] = Bucket(file_item.crop_width, file_item.crop_height)
             self.buckets[bucket_key].file_list_idx.append(idx)
@@ -311,6 +359,15 @@ class BucketsMixin:
         if not quiet:
             print_acc(f'Bucket sizes for {self.dataset_path}:')
             for key, bucket in self.buckets.items():
+                if is_pooled:
+                    # per-source breakdown shows that mixing actually happened
+                    per_source: Dict[str, int] = {}
+                    for i in bucket.file_list_idx:
+                        src = file_list[i].dataset_config.folder_path or file_list[i].dataset_config.dataset_path
+                        per_source[src] = per_source.get(src, 0) + 1
+                    breakdown = ', '.join(f'{os.path.basename(os.path.normpath(k))}={v}' for k, v in per_source.items())
+                    print_acc(f'{key}: {len(bucket.file_list_idx)} files ({breakdown})')
+                else:
                     print_acc(f'{key}: {len(bucket.file_list_idx)} files')
             print_acc(f'{len(self.buckets)} buckets made')
 
@@ -1995,7 +2052,7 @@ class LatentCachingMixin:
                         cached_state_dict = load_file(prep_latent_path, device='cpu') if to_memory else None
                         return prep_item, prep_latent_path, cached_state_dict, False
                     # not saved to disk, load the image/video/audio
-                    prep_item.load_and_process_image(self.transform, only_load_latents=True)
+                    prep_item.load_and_process_image(prep_item.dataloader_transforms or self.transform, only_load_latents=True)
                 except Exception as e:
                     print_acc(f"Error processing image: {prep_item.path}")
                     print_acc(f"Error: {str(e)}")
@@ -2114,7 +2171,7 @@ class LatentCachingMixin:
             frames = None
             # add batch dimension
             cache_uint8 = getattr(self.sd, 'cache_latents_as_uint8', False)
-            if self.dataset_config.cache_tensors_to_disk:
+            if file_item.dataset_config.cache_tensors_to_disk:
                 if not file_item.is_audio_model:
                     tensor_uint8 = _latent_to_uint8(file_item.tensor).cpu()
                     if to_disk:
@@ -2155,7 +2212,7 @@ class LatentCachingMixin:
                 raise e
             # do first frame
             is_video = file_item.is_video
-            if is_video and self.dataset_config.do_i2v:
+            if is_video and file_item.dataset_config.do_i2v:
                 frames = file_item.tensor.unsqueeze(0).to(device, dtype=dtype)
                 if len(frames.shape) == 4:
                     first_frames = frames
@@ -2442,20 +2499,21 @@ class TextEmbeddingCachingMixin:
                 text_embedding_path = file_item.get_text_embedding_path(recalculate=True)
                 # (path, caption) pairs to encode for this item
                 encode_targets = [(text_embedding_path, file_item.caption)]
-                if self.dataset_config.diff_output_preservation:
+                item_config = file_item.dataset_config
+                if item_config.diff_output_preservation:
                     dop_path = file_item.get_dop_text_embedding_path(recalculate=True)
                     if dop_path != text_embedding_path:
                         # trigger word was in the caption, cache the DOP version too
                         encode_targets.append((dop_path, file_item.caption_dop))
                 # dropout embeds are encoded as plain text (no control images)
                 dropout_target_paths = set()
-                if self.dataset_config.caption_dropout_rate > 0:
+                if item_config.caption_dropout_rate > 0:
                     blank_path = file_item.get_blank_text_embedding_path(recalculate=True)
                     if blank_path != text_embedding_path:
                         # cache the dropout caption embedding (blank, or trigger word only)
                         encode_targets.append((blank_path, file_item.get_dropout_caption()))
                         dropout_target_paths.add(blank_path)
-                    if self.dataset_config.diff_output_preservation:
+                    if item_config.diff_output_preservation:
                         # cache the DOP version of the dropout caption (class only)
                         dop_blank_path = file_item.get_dop_blank_text_embedding_path(recalculate=True)
                         if dop_blank_path not in [t[0] for t in encode_targets] + [text_embedding_path]:
@@ -2500,7 +2558,7 @@ class TextEmbeddingCachingMixin:
                         # exactly like its latent rows (frame count / trim)
                         ctrl_img_list.extend(control_video_paths)
                         if len(control_video_paths) > 0:
-                            self.sd._ref_video_dataset_config = self.dataset_config
+                            self.sd._ref_video_dataset_config = item_config
                         
                         if len(ctrl_img_list) == 0:
                             ctrl_img = None
@@ -2528,13 +2586,13 @@ class TextEmbeddingCachingMixin:
                             del prompt_embeds
                     elif (
                         getattr(self.sd, 'encode_first_frame_in_text_embeddings', False)
-                        and self.dataset_config.do_i2v
+                        and item_config.do_i2v
                         and file_item.is_video
                     ):
                         # video item: encode the clip's FIRST FRAME into the text embeddings
                         # as a vision reference, matching sampling (where the ctrl image goes
                         # into the embeds and is held as the clean first frames)
-                        file_item.load_and_process_image(self.transform, only_load_latents=True)
+                        file_item.load_and_process_image(file_item.dataloader_transforms or self.transform, only_load_latents=True)
                         frames = file_item.tensor  # (T, C, H, W) or (C, H, W), in [-1, 1]
                         first = frames[0] if frames.dim() == 4 else frames
                         ctrl_img = (
@@ -2576,7 +2634,7 @@ class TextEmbeddingCachingMixin:
                         file_item.get_dopsd_text_embedding_path(recalculate=True),
                         file_item.caption_dopsd,
                     )]
-                    if self.dataset_config.caption_dropout_rate > 0:
+                    if item_config.caption_dropout_rate > 0:
                         dopsd_blank_path = file_item.get_dopsd_blank_text_embedding_path(recalculate=True)
                         if dopsd_blank_path != dopsd_targets[0][0]:
                             dopsd_targets.append((dopsd_blank_path, file_item.get_dopsd_dropout_caption()))
@@ -2588,10 +2646,10 @@ class TextEmbeddingCachingMixin:
                         if file_item.is_video:
                             # own path rides through the video-ref presentation
                             ctrl_img = [file_item.path]
-                            self.sd._ref_video_dataset_config = self.dataset_config
+                            self.sd._ref_video_dataset_config = item_config
                         else:
                             # own bucketed pixels as the reference image
-                            file_item.load_and_process_image(self.transform, only_load_latents=True)
+                            file_item.load_and_process_image(file_item.dataloader_transforms or self.transform, only_load_latents=True)
                             img = file_item.tensor  # (C, H, W) in [-1, 1]
                             ctrl_img = [
                                 ((img + 1.0) / 2.0)
@@ -2670,7 +2728,11 @@ class CLIPCachingMixin:
 
             # cache unconditionals
             print_acc(f" - Caching {self.clip_vision_num_unconditional_cache} unconditional clip vision to disk")
-            clip_vision_cache_path = os.path.join(self.dataset_config.clip_image_path, '_clip_vision_cache')
+            clip_image_path = next(
+                (c.clip_image_path for c in getattr(self, 'source_configs', [self.dataset_config]) if c.clip_image_path),
+                self.dataset_config.clip_image_path,
+            )
+            clip_vision_cache_path = os.path.join(clip_image_path, '_clip_vision_cache')
 
             unconditional_paths = []
 
@@ -2826,7 +2888,7 @@ class ControlCachingMixin:
 
             # use tqdm to show progress
             for file_item in tqdm(self.file_list, desc=f'Generating Controls'):
-                for control_type in self.dataset_config.controls:
+                for control_type in file_item.dataset_config.controls:
                     # generates the control if it is not already there
                     control_path = self.control_generator.get_control_path(file_item.path, control_type)
                     if control_path is not None:

ui/src/app/jobs/new/SimpleJob.tsx

@@ -1157,6 +1157,17 @@ export default function SimpleJob({
                     </>
                   )}
                 </div>
+                <div>
+                  <Checkbox
+                    label="Pool Datasets"
+                    docKey={'train.pool_datasets'}
+                    className="pt-1"
+                    checked={jobConfig.config.process[0].train.pool_datasets || false}
+                    onChange={value => {
+                      setJobConfig(value ? true : undefined, 'config.process[0].train.pool_datasets');
+                    }}
+                  />
+                </div>
               </div>
             </Card>
           )}

ui/src/docs.tsx

@@ -288,6 +288,30 @@ const docs: { [key: string]: ConfigDoc } = {
       </>
     ),
   },
+  'train.pool_datasets': {
+    title: 'Pool Datasets',
+    description: (
+      <>
+        Off by default. Only useful when you have numerous datasets.
+        <br />
+        <br />
+        Normally every dataset is loaded as its own unit with its own buckets, so each batch is built from a single
+        dataset and each dataset keeps its own loader threads. With many datasets this means batches never mix images
+        from different folders and the number of loader threads grows with the dataset count.
+        <br />
+        <br />
+        Pooling merges all of your datasets into one shared set of buckets. Images from different folders that land in
+        the same bucket are mixed together in the same batch, and only one set of loader threads is used for the whole
+        pool. Regularization datasets are pooled separately from your training datasets, never together.
+        <br />
+        <br />
+        Each dataset keeps its own settings (captions, control images, repeats, flips, scale, resolution). Datasets whose
+        settings would not collate into one batch (for example a different number of control images, video vs image,
+        or different caching) are automatically kept in separate buckets within the pool. All pooled datasets must use
+        the same batch size.
+      </>
+    ),
+  },
   'train.do_differential_guidance': {
     title: 'Differential Guidance',
     description: (

ui/src/types.ts

@@ -187,6 +187,7 @@ export interface TrainConfig {
   loss_type: 'mse' | 'mae' | 'wavelet' | 'stepped';
   do_differential_guidance?: boolean;
   differential_guidance_scale?: number;
+  pool_datasets?: boolean;
   audio_loss_multiplier?: number;
   max_loss?: number | null;
   validation_config?: ValidationConfig;