From 1fd6666682017bf4737025d58003746bcab09d48 Mon Sep 17 00:00:00 2001
From: psychedelicious <4822129+psychedelicious@users.noreply.github.com>
Date: Thu, 30 Nov 2023 19:57:29 +1100
Subject: [PATCH] feat(backend): clear latents files on startup

Adds logic to `DiskLatentsStorage.start()` to empty the latents folder on startup.

Adds start and stop methods to `ForwardCacheLatentsStorage`. This is required for `DiskLatentsStorage.start()` to be called, due to how this particular service breaks the direct DI pattern, wrapping the underlying storage with a cache.
---
 .../latents_storage/latents_storage_disk.py   | 19 +++++++++++++++++++
 .../latents_storage_forward_cache.py          | 14 ++++++++++++++
 2 files changed, 33 insertions(+)

diff --git a/invokeai/app/services/latents_storage/latents_storage_disk.py b/invokeai/app/services/latents_storage/latents_storage_disk.py
index 6e7010bae0..58d1df5081 100644
--- a/invokeai/app/services/latents_storage/latents_storage_disk.py
+++ b/invokeai/app/services/latents_storage/latents_storage_disk.py
@@ -5,6 +5,8 @@ from typing import Union
 
 import torch
 
+from invokeai.app.services.invoker import Invoker
+
 from .latents_storage_base import LatentsStorageBase
 
 
@@ -17,6 +19,23 @@ class DiskLatentsStorage(LatentsStorageBase):
         self.__output_folder = output_folder if isinstance(output_folder, Path) else Path(output_folder)
         self.__output_folder.mkdir(parents=True, exist_ok=True)
 
+    def start(self, invoker: Invoker) -> None:
+        self._invoker = invoker
+
+        # Delete all latents files on startup
+        deleted_latents_count = 0
+        freed_space = 0
+        for latents_file in Path(self.__output_folder).glob("*"):
+            if latents_file.is_file():
+                freed_space += latents_file.stat().st_size
+                deleted_latents_count += 1
+                latents_file.unlink()
+        if deleted_latents_count > 0:
+            freed_space_in_mb = round(freed_space / 1024 / 1024, 2)
+            self._invoker.services.logger.info(
+                f"Deleted {deleted_latents_count} latents files, freeing {freed_space_in_mb}MB"
+            )
+
     def get(self, name: str) -> torch.Tensor:
         latent_path = self.get_path(name)
         return torch.load(latent_path)
diff --git a/invokeai/app/services/latents_storage/latents_storage_forward_cache.py b/invokeai/app/services/latents_storage/latents_storage_forward_cache.py
index da82b5904d..6232b76a27 100644
--- a/invokeai/app/services/latents_storage/latents_storage_forward_cache.py
+++ b/invokeai/app/services/latents_storage/latents_storage_forward_cache.py
@@ -5,6 +5,8 @@ from typing import Dict, Optional
 
 import torch
 
+from invokeai.app.services.invoker import Invoker
+
 from .latents_storage_base import LatentsStorageBase
 
 
@@ -23,6 +25,18 @@ class ForwardCacheLatentsStorage(LatentsStorageBase):
         self.__cache_ids = Queue()
         self.__max_cache_size = max_cache_size
 
+    def start(self, invoker: Invoker) -> None:
+        self._invoker = invoker
+        start_op = getattr(self.__underlying_storage, "start", None)
+        if callable(start_op):
+            start_op(invoker)
+
+    def stop(self, invoker: Invoker) -> None:
+        self._invoker = invoker
+        stop_op = getattr(self.__underlying_storage, "stop", None)
+        if callable(stop_op):
+            stop_op(invoker)
+
     def get(self, name: str) -> torch.Tensor:
         cache_item = self.__get_cache(name)
         if cache_item is not None: