Add threadding to bucket resizing of a batch. Helpful to speed up large batch gathering.

toolkit/config_modules.py

@@ -1055,6 +1055,9 @@ class DatasetConfig:
         # which can cause severe PCIe thrashing for users at the VRAM ceiling.
         # Opt in if you have stable VRAM headroom and want the transfer speedup.
         self.pin_memory: bool = kwargs.get('pin_memory', False)
+        # threads that decode/resize the images of one bucketed batch concurrently inside a
+        # dataloader worker; 0 or 1 loads the batch serially
+        self.batch_load_threads: int = kwargs.get('batch_load_threads', min(8, os.cpu_count() or 1))
         # threads used to prep (decode/resize) items ahead of the VAE while caching latents
         self.cache_latents_num_workers: int = kwargs.get('cache_latents_num_workers', min(6, os.cpu_count() or 1))
         self.extra_values: List[float] = kwargs.get('extra_values', [])

toolkit/data_loader.py

@@ -3,6 +3,7 @@ import json
 import os
 import random
 import traceback
+from concurrent.futures import ThreadPoolExecutor
 from functools import lru_cache
 from typing import List, TYPE_CHECKING
 
@@ -420,6 +421,8 @@ class AiToolkitDataset(LatentCachingMixin, ControlCachingMixin, CLIPCachingMixin
         self.random_scale = dataset_config.random_scale
         self.scale = dataset_config.scale
         self.batch_size = batch_size
+        # built lazily inside the worker; an executor cannot be pickled to spawned workers
+        self._batch_pool = None
         # we always random crop if random scale is enabled
         self.random_crop = self.random_scale if self.random_scale else dataset_config.random_crop
         self.resolution = dataset_config.resolution
@@ -656,7 +659,12 @@ class AiToolkitDataset(LatentCachingMixin, ControlCachingMixin, CLIPCachingMixin
                 # tried everything to solve this. No way to reset length when redoing things. Pick another index
                 item = random.randint(0, len(self.batch_indices) - 1)
             idx_list = self.batch_indices[item]
+            n_threads = min(self.dataset_config.batch_load_threads, len(idx_list))
+            if n_threads <= 1:
                 return [self._get_single_item(idx) for idx in idx_list]
+            if self._batch_pool is None or self._batch_pool._max_workers != n_threads:
+                self._batch_pool = ThreadPoolExecutor(max_workers=n_threads)
+            return list(self._batch_pool.map(self._get_single_item, idx_list))
         else:
             # Dataloader is batching
             return self._get_single_item(item)