From 3d472de2f13f7c9b370f9fa3f96b86eb2776270a Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Wed, 29 Jul 2026 17:47:04 -0600 Subject: [PATCH 1/2] Reworked spawning on dataloader so windows and mac can use multiple data loader workers now --- toolkit/data_loader.py | 43 ++++++++++++++++++++++++++++++------------ 1 file changed, 31 insertions(+), 12 deletions(-) diff --git a/toolkit/data_loader.py b/toolkit/data_loader.py index ee90e263..f8efbace 100644 --- a/toolkit/data_loader.py +++ b/toolkit/data_loader.py @@ -606,6 +606,14 @@ class AiToolkitDataset(LatentCachingMixin, ControlCachingMixin, CLIPCachingMixin self.setup_controls() self.epoch_num += 1 + def __getstate__(self): + # on Windows/macOS dataloader workers are spawned, which pickles the dataset. + # sd (the model) is not picklable (weakrefs, cuda tensors) and is only needed + # for caching, which runs in the main process before iteration starts. + state = self.__dict__.copy() + state['sd'] = None + return state + def __len__(self): if self.dataset_config.buckets: return len(self.batch_indices) @@ -652,6 +660,13 @@ class AiToolkitDataset(LatentCachingMixin, ControlCachingMixin, CLIPCachingMixin return self._get_single_item(item) +def dto_collation(batch: List['FileItemDTO']): + # must be a module level function so spawned dataloader workers can pickle it + return DataLoaderBatchDTO( + file_items=batch + ) + + def get_dataloader_from_datasets( dataset_options, batch_size=1, @@ -692,22 +707,26 @@ def get_dataloader_from_datasets( # todo build scheduler that can get buckets from all datasets that match # todo and evenly distribute reg images - def dto_collation(batch: List['FileItemDTO']): - # create DTO batch - batch = DataLoaderBatchDTO( - file_items=batch - ) - return batch - # check if is caching latents dataloader_kwargs = {} - - if is_native_windows() or is_macos(): - dataloader_kwargs['num_workers'] = 0 - else: - dataloader_kwargs['num_workers'] = dataset_config_list[0].num_workers + + dataloader_kwargs['num_workers'] = dataset_config_list[0].num_workers + if dataloader_kwargs['num_workers'] > 0: dataloader_kwargs['prefetch_factor'] = dataset_config_list[0].prefetch_factor + # keep workers alive across epochs. Without this, spawn platforms (Windows/macOS) + # boot new worker processes every epoch, which can take longer than the epoch + # itself on small datasets. The dataset is static after epoch 0 (setup_epoch only + # does work on the first call) and per-epoch shuffling happens in the main process + # sampler, so workers never hold stale state. + dataloader_kwargs['persistent_workers'] = True + # spawned workers re-import the full stack at boot and would repeat every + # import-time warning the parent already printed. Children inherit these env + # vars; the parent is unaffected since its imports already happened. + os.environ.setdefault('PYTHONWARNINGS', 'ignore::FutureWarning') + os.environ.setdefault('TORCH_LOGS', '-torch.utils._pytree') + os.environ.setdefault('DIFFUSERS_VERBOSITY', 'error') + os.environ.setdefault('NO_ALBUMENTATIONS_UPDATE', '1') if has_buckets: # make sure they all have buckets From 23f1ebfb763b38e078fb64a13e38d559988a78d5 Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Wed, 29 Jul 2026 17:47:53 -0600 Subject: [PATCH 2/2] Update to high speed xet env var --- run.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/run.py b/run.py index 3ee5708f..a48e16a9 100644 --- a/run.py +++ b/run.py @@ -3,7 +3,7 @@ import sys from dotenv import load_dotenv # Load the .env file if it exists load_dotenv() -os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = os.getenv("HF_HUB_ENABLE_HF_TRANSFER", "1") +os.environ["HF_XET_HIGH_PERFORMANCE"] = os.getenv("HF_XET_HIGH_PERFORMANCE", "1") os.environ["HF_HUB_DISABLE_XET"] = os.getenv("HF_HUB_DISABLE_XET", "0") os.environ["NO_ALBUMENTATIONS_UPDATE"] = "1" seed = None