diff --git a/toolkit/config_modules.py b/toolkit/config_modules.py index cda8662..9de7934 100644 --- a/toolkit/config_modules.py +++ b/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 diff --git a/toolkit/dataloader_mixins.py b/toolkit/dataloader_mixins.py index 42b1440..02c7afa 100644 --- a/toolkit/dataloader_mixins.py +++ b/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,111 +1814,154 @@ 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 ""}'): - file_item.is_caching_to_disk = to_disk - file_item.is_caching_to_memory = to_memory - file_item.latent_load_device = self.sd.device - - latent_path = file_item.get_latent_path(recalculate=True) - # check if it is saved to disk already - if os.path.exists(latent_path): - if to_memory: - # load it into memory - state_dict = load_file(latent_path, device='cpu') - cached_latent = state_dict['latent'] - if cached_latent.dtype == torch.uint8: - # pixel-space latents cached as uint8 - cached_latent = _latent_from_uint8(cached_latent) - file_item._encoded_latent = cached_latent.to('cpu', dtype=self.sd.torch_dtype) - if 'first_frame_latent' in state_dict: - cached_first_frame = state_dict['first_frame_latent'] - if cached_first_frame.dtype == torch.uint8: - cached_first_frame = _latent_from_uint8(cached_first_frame) - file_item._cached_first_frame_latent = cached_first_frame.to('cpu', dtype=self.sd.torch_dtype) - if 'audio_latent' in state_dict: - 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 - 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() - first_frame_latent = None - audio_latent = None - frames = None - # add batch dimension - cache_uint8 = getattr(self.sd, 'cache_latents_as_uint8', False) - try: - imgs = file_item.tensor.unsqueeze(0).to(device, dtype=dtype) - latent = self.sd.encode_images(imgs).squeeze(0) - if to_disk: - if cache_uint8: - state_dict['latent'] = _latent_to_uint8(latent).cpu() - else: - state_dict['latent'] = latent.clone().detach().cpu() - except Exception as e: - print_acc(f"Error processing image: {file_item.path}") - print_acc(f"Error: {str(e)}") - raise e - # do first frame - is_video = self.dataset_config.auto_frame_count or self.dataset_config.num_frames > 1 - if is_video and self.dataset_config.do_i2v: - frames = file_item.tensor.unsqueeze(0).to(device, dtype=dtype) - if len(frames.shape) == 4: - first_frames = frames - elif len(frames.shape) == 5: - first_frames = frames[:, 0] - else: - raise ValueError(f"Unknown frame shape {frames.shape}") - first_frame_latent = self.sd.encode_images(first_frames).squeeze(0) - if to_disk: - if cache_uint8: - state_dict['first_frame_latent'] = _latent_to_uint8(first_frame_latent).cpu() - else: - state_dict['first_frame_latent'] = first_frame_latent.clone().detach().cpu() - - # audio (video+audio models only — audio-only models already encoded above via encode_images) - if not self.is_audio_model and file_item.audio_data is not None: - audio_latent = self.sd.encode_audio([file_item.audio_data]).squeeze(0) - if to_disk: - state_dict['audio_latent'] = audio_latent.clone().detach().cpu() - - if is_video: - state_dict['num_frames'] = torch.tensor(file_item.num_frames, dtype=torch.int32) - - # save_latent - if to_disk: - # metadata - meta = get_meta_for_safetensors(file_item.get_latent_info_dict()) - os.makedirs(os.path.dirname(latent_path), exist_ok=True) - save_file(state_dict, latent_path, metadata=meta) - - if to_memory: - # keep it in memory - file_item._encoded_latent = latent.to('cpu', dtype=self.sd.torch_dtype) - if first_frame_latent is not None: - file_item._cached_first_frame_latent = first_frame_latent.to('cpu', dtype=self.sd.torch_dtype) - if audio_latent is not None: - file_item._cached_audio_latent = audio_latent.to('cpu', dtype=self.sd.torch_dtype) - - del imgs - del latent - del frames - del file_item.tensor - del state_dict - del first_frame_latent - del audio_latent - file_item.cleanup() - - file_item.is_latent_cached = True - i += 1 + pbar = tqdm(total=len(self.file_list), desc=f'Caching latents{" to disk" if to_disk else ""}') + executor = ThreadPoolExecutor(max_workers=num_workers) + try: + 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() # 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 not needs_encode: + if to_memory: + # load it into memory + state_dict = cached_state_dict + cached_latent = state_dict['latent'] + if cached_latent.dtype == torch.uint8: + # pixel-space latents cached as uint8 + cached_latent = _latent_from_uint8(cached_latent) + file_item._encoded_latent = cached_latent.to('cpu', dtype=self.sd.torch_dtype) + if 'first_frame_latent' in state_dict: + cached_first_frame = state_dict['first_frame_latent'] + if cached_first_frame.dtype == torch.uint8: + cached_first_frame = _latent_from_uint8(cached_first_frame) + file_item._cached_first_frame_latent = cached_first_frame.to('cpu', dtype=self.sd.torch_dtype) + if 'audio_latent' in state_dict: + file_item._cached_audio_latent = state_dict['audio_latent'].to('cpu', dtype=self.sd.torch_dtype) + else: + # not saved to disk, calculate + # the image/video/audio was already loaded by the prep thread + dtype = self.sd.torch_dtype + device = self.sd.device_torch + state_dict = OrderedDict() + first_frame_latent = None + audio_latent = None + frames = None + # add batch dimension + cache_uint8 = getattr(self.sd, 'cache_latents_as_uint8', False) + try: + imgs = file_item.tensor.unsqueeze(0).to(device, dtype=dtype) + latent = self.sd.encode_images(imgs).squeeze(0) + if to_disk: + if cache_uint8: + state_dict['latent'] = _latent_to_uint8(latent).cpu() + else: + state_dict['latent'] = latent.clone().detach().cpu() + except Exception as e: + print_acc(f"Error processing image: {file_item.path}") + print_acc(f"Error: {str(e)}") + raise e + # do first frame + is_video = self.dataset_config.auto_frame_count or self.dataset_config.num_frames > 1 + if is_video and self.dataset_config.do_i2v: + frames = file_item.tensor.unsqueeze(0).to(device, dtype=dtype) + if len(frames.shape) == 4: + first_frames = frames + elif len(frames.shape) == 5: + first_frames = frames[:, 0] + else: + raise ValueError(f"Unknown frame shape {frames.shape}") + first_frame_latent = self.sd.encode_images(first_frames).squeeze(0) + if to_disk: + if cache_uint8: + state_dict['first_frame_latent'] = _latent_to_uint8(first_frame_latent).cpu() + else: + state_dict['first_frame_latent'] = first_frame_latent.clone().detach().cpu() + + # audio (video+audio models only — audio-only models already encoded above via encode_images) + if not self.is_audio_model and file_item.audio_data is not None: + audio_latent = self.sd.encode_audio([file_item.audio_data]).squeeze(0) + if to_disk: + state_dict['audio_latent'] = audio_latent.clone().detach().cpu() + + if is_video: + state_dict['num_frames'] = torch.tensor(file_item.num_frames, dtype=torch.int32) + + # save_latent + if to_disk: + # metadata + meta = get_meta_for_safetensors(file_item.get_latent_info_dict()) + os.makedirs(os.path.dirname(latent_path), exist_ok=True) + save_file(state_dict, latent_path, metadata=meta) + + if to_memory: + # keep it in memory + file_item._encoded_latent = latent.to('cpu', dtype=self.sd.torch_dtype) + if first_frame_latent is not None: + file_item._cached_first_frame_latent = first_frame_latent.to('cpu', dtype=self.sd.torch_dtype) + if audio_latent is not None: + file_item._cached_audio_latent = audio_latent.to('cpu', dtype=self.sd.torch_dtype) + + del imgs + del latent + del frames + del file_item.tensor + del state_dict + del first_frame_latent + del audio_latent + file_item.cleanup() + class TextEmbeddingFileItemDTOMixin: def __init__(self, *args, **kwargs):