Queue up videos with multiple threads when caching latents so the VAE is not waiting on videos to process

toolkit/config_modules.py

@@ -1030,6 +1030,8 @@ class DatasetConfig:
 
         self.num_workers: int = kwargs.get('num_workers', 2)
         self.prefetch_factor: int = kwargs.get('prefetch_factor', 2)
+        # 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', [])
         self.square_crop: bool = kwargs.get('square_crop', False)
         # apply same augmentations to control images. Usually want this true unless special case

toolkit/dataloader_mixins.py

@@ -1,11 +1,13 @@
 import base64
 import glob
 import hashlib
+import itertools
 import json
 import math
 import os
 import random
-from collections import OrderedDict
+from collections import OrderedDict, deque
+from concurrent.futures import ThreadPoolExecutor
 from typing import TYPE_CHECKING, List, Dict, Union
 import traceback
 
@@ -1812,19 +1814,69 @@ class LatentCachingMixin:
             # move sd items to cpu except for vae
             self.sd.set_device_state_preset('cache_latents')
 
+            # prep (video decode, frame extraction, audio load, disk reads) is done by a
+            # thread pool so the next items are ready while the current one is encoding.
+            # the in-flight window is bounded so decoded videos don't pile up in RAM.
+            num_workers = max(1, self.dataset_config.cache_latents_num_workers)
+
+            def _prep(prep_item: 'FileItemDTO'):
+                prep_item.is_caching_to_disk = to_disk
+                prep_item.is_caching_to_memory = to_memory
+                prep_item.latent_load_device = self.sd.device
+
+                prep_latent_path = prep_item.get_latent_path(recalculate=True)
+                try:
+                    if os.path.exists(prep_latent_path):
+                        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)
+                except Exception as e:
+                    print_acc(f"Error processing image: {prep_item.path}")
+                    print_acc(f"Error: {str(e)}")
+                    raise e
+                return prep_item, prep_latent_path, None, True
+
             # use tqdm to show progress
             i = 0
-            for file_item in tqdm(self.file_list, desc=f'Caching latents{" to disk" if to_disk else ""}'):
+            pbar = tqdm(total=len(self.file_list), desc=f'Caching latents{" to disk" if to_disk else ""}')
-                file_item.is_caching_to_disk = to_disk
+            executor = ThreadPoolExecutor(max_workers=num_workers)
-                file_item.is_caching_to_memory = to_memory
+            try:
-                file_item.latent_load_device = self.sd.device
+                pending = deque()
+                file_iter = iter(self.file_list)
+                for queued_item in itertools.islice(file_iter, num_workers + 2):
+                    pending.append(executor.submit(_prep, queued_item))
+                while pending:
+                    file_item, latent_path, cached_state_dict, needs_encode = pending.popleft().result()
+                    # keep the window full
+                    next_item = next(file_iter, None)
+                    if next_item is not None:
+                        pending.append(executor.submit(_prep, next_item))
+                    self._cache_one_latent(file_item, latent_path, cached_state_dict, needs_encode, to_disk, to_memory)
+                    file_item.is_latent_cached = True
+                    i += 1
+                    pbar.update(1)
+            finally:
+                executor.shutdown(wait=True, cancel_futures=True)
+                pbar.close()
 
-                latent_path = file_item.get_latent_path(recalculate=True)
+            # restore device state
+            self.sd.restore_device_state()
+
+    def _cache_one_latent(
+            self: 'AiToolkitDataset',
+            file_item: 'FileItemDTO',
+            latent_path: str,
+            cached_state_dict,
+            needs_encode: bool,
+            to_disk: bool,
+            to_memory: bool,
+    ):
         # check if it is saved to disk already
-                if os.path.exists(latent_path):
+        if not needs_encode:
             if to_memory:
                 # load it into memory
-                        state_dict = load_file(latent_path, device='cpu')
+                state_dict = cached_state_dict
                 cached_latent = state_dict['latent']
                 if cached_latent.dtype == torch.uint8:
                     # pixel-space latents cached as uint8
@@ -1839,8 +1891,7 @@ class LatentCachingMixin:
                     file_item._cached_audio_latent = state_dict['audio_latent'].to('cpu', dtype=self.sd.torch_dtype)
         else:
             # not saved to disk, calculate
-                    # load the image first
+            # the image/video/audio was already loaded by the prep thread
-                    file_item.load_and_process_image(self.transform, only_load_latents=True)
             dtype = self.sd.torch_dtype
             device = self.sd.device_torch
             state_dict = OrderedDict()
@@ -1911,12 +1962,6 @@ class LatentCachingMixin:
             del audio_latent
             file_item.cleanup()
 
-                file_item.is_latent_cached = True
-                i += 1
-
-            # restore device state
-            self.sd.restore_device_state()
-
 
 class TextEmbeddingFileItemDTOMixin:
     def __init__(self, *args, **kwargs):