From 49479cdaa6804ad0b41e03479f124a03624b127f Mon Sep 17 00:00:00 2001
From: asagi4 <130366179+asagi4@users.noreply.github.com>
Date: Thu, 8 Feb 2024 17:46:25 +0200
Subject: [PATCH] Add a node that runs torch / python gc

---
 utils/memory.py | 23 +++++++++++++++++++++++
 1 file changed, 23 insertions(+)
 create mode 100644 utils/memory.py

diff --git a/utils/memory.py b/utils/memory.py
new file mode 100644
index 0000000..f400788
--- /dev/null
+++ b/utils/memory.py
@@ -0,0 +1,23 @@
+import gc
+import comfy.model_management
+import logging
+
+log = logging.getLogger("comfyui-misc-utils")
+
+
+class MUForceCacheClear:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required": {"model": ("MODEL",)}}
+
+    NOT_IDEMPOTENT = True
+    RETURN_TYPES = ("MODEL",)
+    FUNCTION = "do"
+
+    CATEGORY = "utils"
+
+    def do(self, model):
+        log.info("Running Python GC and clearing CUDA cache")
+        gc.collect()
+        comfy.model_management.soft_empty_cache(True)
+        return (model,)