From add525c7c08364ddb2322ff5f78995d06e3ef7c6 Mon Sep 17 00:00:00 2001
From: silveroxides <ishimarukaito@gmail.com>
Date: Tue, 22 Jul 2025 20:07:21 +0200
Subject: [PATCH] prep for adding keyhandler

---
 __init__.py            |  2 ++
 nodes/__init__.py      |  2 ++
 nodes/keyhandler.py    |  0
 nodes/modelunloader.py | 31 +++++++++++++++++++++++++++++++
 4 files changed, 35 insertions(+)
 create mode 100644 nodes/keyhandler.py
 create mode 100644 nodes/modelunloader.py

diff --git a/__init__.py b/__init__.py
index bf2efef..6604a8e 100644
--- a/__init__.py
+++ b/__init__.py
@@ -5,6 +5,7 @@ NODE_CLASS_MAPPINGS = {
     "CLIPMetaKeys": CLIPMetaKeys,
     "LoRAMetaKeys": LoRAMetaKeys,
     "CheckpointMetaKeys": CheckpointMetaKeys,
+    "OffloadModels": OffloadModelsNode,
 }
 
 NODE_DISPLAY_NAME_MAPPINGS = {
@@ -12,4 +13,5 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "CLIPMetaKeys": "Get CLIP Metadata and Keys",
     "LoRAMetaKeys": "Get LoRA Metadata and Keys",
     "CheckpointMetaKeys": "Get Checkpoint Metadata and Keys",
+    "OffloadModelsNode": "Offload Models with Passthrough",
 }
\ No newline at end of file
diff --git a/nodes/__init__.py b/nodes/__init__.py
index c59fb29..c673cc3 100644
--- a/nodes/__init__.py
+++ b/nodes/__init__.py
@@ -1,8 +1,10 @@
 from .metakeys import *
+from .modelunloader import *
 
 __all__ = (
     "UNetMetaKeys",
     "CLIPMetaKeys",
     "LoRAMetaKeys",
     "CheckpointMetaKeys",
+    "OffloadModels",
 )
\ No newline at end of file
diff --git a/nodes/keyhandler.py b/nodes/keyhandler.py
new file mode 100644
index 0000000..e69de29
diff --git a/nodes/modelunloader.py b/nodes/modelunloader.py
new file mode 100644
index 0000000..81d8881
--- /dev/null
+++ b/nodes/modelunloader.py
@@ -0,0 +1,31 @@
+from comfy.comfy_types import IO, ComfyNodeABC, InputTypeDict
+import comfy.model_management as mm
+
+class OffloadModelsNode(ComfyNodeABC):
+    @classmethod
+    def INPUT_TYPES(s) -> InputTypeDict:
+        return {
+            "required": {
+                "any": (IO.ANY, {}),
+                "force": (IO.BOOLEAN, {"default": False, "tooltip": "If True, forcibly unloads models from GPU."}),
+                "enable_offload": (IO.BOOLEAN, {"default": False, "tooltip": "If True, calls ComfyUI's `unload_all_models()` function."}),
+            }
+        }
+
+    RETURN_TYPES = (IO.ANY, )
+    RETURN_NAMES = ("any", )
+    FUNCTION = 'function'
+    CATEGORY = 'advanced'
+
+
+    def function(self, any, force, enable_offload):
+        if enable_offload == True:
+            print("OffloadModels: `enable_offload` is True. Calling `mm.unload_all_models()`.")
+            mm.unload_all_models()
+            if Force == True:
+                mm.soft_empty_cache(Force=True)
+            else:
+                mm.soft_empty_cache()
+        return (any,)
+
+__all__ = ("OffloadModelsNode")
\ No newline at end of file