From 546313b439835a0f6c410246de7e3776bd29f767 Mon Sep 17 00:00:00 2001
From: kijai <40791699+kijai@users.noreply.github.com>
Date: Mon, 16 Oct 2023 02:31:04 +0300
Subject: [PATCH] Add prototype of VRAM debug node

Returns free VRAM as int, both before and after the cleaning operations the node also tries to do.  Currently only has input for clip_vision which seems to fail to unload otherwise (when used with IPAdapter)
---
 nodes.py | 36 +++++++++++++++++++++++++++++++++++-
 1 file changed, 35 insertions(+), 1 deletion(-)

diff --git a/nodes.py b/nodes.py
index cd1a7ef..d6e43e5 100644
--- a/nodes.py
+++ b/nodes.py
@@ -8,6 +8,7 @@ import os
 import librosa
 from scipy.special import erf
 from .fluid import Fluid
+import comfy.model_management
 
 from nodes import MAX_RESOLUTION
 
@@ -632,7 +633,38 @@ class ConditioningSetMaskAndCombine4:
         for t in negative_4:
             append_helper(t, mask_4, c2, set_area_to_bounds, strength)
         return (c, c2)
-     
+    
+class VRAM_Debug:
+    
+    @classmethod
+    
+    def INPUT_TYPES(s):
+      return {
+        "required": {
+			  "model": ("MODEL",),
+              
+		  },
+        "optional": {
+            "clip_vision": ("CLIP_VISION", ),
+        }
+	  }
+    RETURN_TYPES = ("MODEL", "INT", "INT",)
+    RETURN_NAMES = ("model", "freemem_before", "freemem_after")
+    FUNCTION = "VRAMdebug"
+    CATEGORY = "KJNodes"
+
+    def VRAMdebug(self, model, clip_vision=None):
+        freemem_before = comfy.model_management.get_free_memory()
+        print(freemem_before)
+        torch.cuda.empty_cache()
+        torch.cuda.ipc_collect()
+        if clip_vision is not None:
+            print("unloading clip_vision_clone")
+            comfy.model_management.unload_model_clones(clip_vision.patcher)
+        freemem_after = comfy.model_management.get_free_memory()
+        print(freemem_after)
+        return (model, freemem_before, freemem_after)
+          
 NODE_CLASS_MAPPINGS = {
     "INTConstant": INTConstant,
     "ConditioningMultiCombine": ConditioningMultiCombine,
@@ -646,6 +678,7 @@ NODE_CLASS_MAPPINGS = {
     "CreateAudioMask": CreateAudioMask,
     "CreateFadeMask": CreateFadeMask,
     "CreateFluidMask" :CreateFluidMask,
+    "VRAM_Debug" : VRAM_Debug,
 }
 NODE_DISPLAY_NAME_MAPPINGS = {
     "INTConstant": "INT Constant",
@@ -659,4 +692,5 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "CreateTextMask" : "CreateTextMask",
     "CreateFadeMask" : "CreateFadeMask",
     "CreateFluidMask" : "CreateFluidMask",
+    "VRAM_Debug" : "VRAM Debug",
 }
\ No newline at end of file