From b5a2669b740de52947c23886cb9b84fc3bafe3f1 Mon Sep 17 00:00:00 2001
From: Jaret Burkett <jaretburkett@gmail.com>
Date: Tue, 12 Sep 2023 07:03:10 -0600
Subject: [PATCH] Fixed memory leak

---
 toolkit/dataloader_mixins.py | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

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()