prep for adding keyhandler

__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",
 }

nodes/__init__.py

@@ -1,8 +1,10 @@
 from .metakeys import *
+from .modelunloader import *
 
 __all__ = (
     "UNetMetaKeys",
     "CLIPMetaKeys",
     "LoRAMetaKeys",
     "CheckpointMetaKeys",
+    "OffloadModels",
 )

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