From 29b4e81030c5653b671c78d967482e48b1fa1bf2 Mon Sep 17 00:00:00 2001
From: tazlin <tazlin.on.github@gmail.com>
Date: Sun, 25 Aug 2024 09:27:01 -0400
Subject: [PATCH] feat: `aggressive_unloading` for more regular garbage
 collection

ComfyUI's default mode of running (`main.py`) includes calls to `cleanup_models(...)` and `soft_empty_cache(...)` on a timer. I suspect issues that have arisen lately are rooted in the fact that horde-engine does not currently do something similar. I am adding this (default on) option to the `HordeLib` class to call these after every pipeline run.
---
 hordelib/comfy_horde.py | 8 ++++++++
 hordelib/horde.py       | 2 ++
 2 files changed, 10 insertions(+)

diff --git a/hordelib/comfy_horde.py b/hordelib/comfy_horde.py
index bf7b95f5..6c913a35 100644
--- a/hordelib/comfy_horde.py
+++ b/hordelib/comfy_horde.py
@@ -427,6 +427,7 @@ def __init__(
         self,
         *,
         comfyui_callback: typing.Callable[[str, dict, str], None] | None = None,
+        aggressive_unloading: bool = True,
     ) -> None:
         """Initialise the Comfy_Horde object.
 
@@ -453,6 +454,7 @@ def __init__(
         self._load_custom_nodes()
 
         self._comfyui_callback = comfyui_callback
+        self.aggressive_unloading = aggressive_unloading
 
     def _set_comfyui_paths(self) -> None:
         # These set the default paths for comfyui to look for models and embeddings. From within hordelib,
@@ -829,6 +831,12 @@ def _run_pipeline(
                     inference.execute(pipeline, self.client_id, {"client_id": self.client_id}, valid[2])
             except Exception as e:
                 logger.exception(f"Exception during comfy execute: {e}")
+            finally:
+                if self.aggressive_unloading:
+                    global _comfy_cleanup_models
+                    logger.debug("Cleaning up models")
+                    _comfy_cleanup_models(False)
+                    _comfy_soft_empty_cache(True)
 
         stdio.replay()
 
diff --git a/hordelib/horde.py b/hordelib/horde.py
index f6e2b088..dee75944 100644
--- a/hordelib/horde.py
+++ b/hordelib/horde.py
@@ -352,10 +352,12 @@ def __init__(
         self,
         *,
         comfyui_callback: Callable[[str, dict, str], None] | None = None,
+        aggressive_unloading: bool = False,
     ):
         if not self._initialised:
             self.generator = Comfy_Horde(
                 comfyui_callback=comfyui_callback if comfyui_callback else self._comfyui_callback,
+                aggressive_unloading=aggressive_unloading,
             )
             self.__class__._initialised = True