diff --git a/toolkit/config_modules.py b/toolkit/config_modules.py index c39a4a9..70dd08a 100644 --- a/toolkit/config_modules.py +++ b/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', []) diff --git a/toolkit/data_loader.py b/toolkit/data_loader.py index 3f96183..a39f586 100644 --- a/toolkit/data_loader.py +++ b/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] - return [self._get_single_item(idx) for idx in idx_list] + 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)