...

clear_vram.py.txt

@@ -0,0 +1,33 @@
+import torch
+import gc
+
+class ClearVRAM:
+    def __init__(self):
+        self.cleared = False
+    
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "clear": ("BOOLEAN", {"default": False, "label": "Clear VRAM"})
+            }
+        }
+    
+    RETURN_TYPES = ("STRING",)
+    FUNCTION = "clear_vram"
+    CATEGORY = "memory_management"
+    OUTPUT_NODE = True
+
+    def clear_vram(self, clear):
+        if clear and not self.cleared:
+            torch.cuda.empty_cache()
+            gc.collect()
+            torch.cuda.synchronize()
+            print("VRAM cleared")
+            self.cleared = True
+            return ("VRAM cleared",)
+        elif not clear:
+            self.cleared = False
+            return ("Ready to clear VRAM",)
+        else:
+            return ("VRAM already cleared",)