diff --git a/toolkit/dataloader_mixins.py b/toolkit/dataloader_mixins.py index 018c3638..09a5bf65 100644 --- a/toolkit/dataloader_mixins.py +++ b/toolkit/dataloader_mixins.py @@ -370,7 +370,7 @@ class LatentCachingMixin: self.sd.set_device_state_preset('cache_latents') # use tqdm to show progress - for file_item in tqdm(self.file_list, desc=f'Caching latents{" to disk" if to_disk else ""}'): + for i, file_item in tqdm(enumerate(self.file_list), desc=f'Caching latents{" to disk" if to_disk else ""}'): # set latent space version if self.sd.is_xl: file_item.latent_space_version = 'sdxl' @@ -410,8 +410,15 @@ class LatentCachingMixin: # keep it in memory file_item._encoded_latent = latent.to('cpu', dtype=self.sd.torch_dtype) + del imgs + del latent + del file_item.tensor + flush(garbage_collect=False) file_item.is_latent_cached = True + # flush every 100 + # if i % 100 == 0: + # flush() # restore device state self.sd.restore_device_state()