From a8c216c7e9c536f7daf3e13a01248a4d3b4911e4 Mon Sep 17 00:00:00 2001
From: kijai <40791699+kijai@users.noreply.github.com>
Date: Sat, 4 May 2024 10:09:17 +0300
Subject: [PATCH] Make VRAM_Debug report the values within the node itself

---
 nodes/nodes.py    | 13 ++++++++-----
 web/js/jsnodes.js | 17 +++++++++++++++++
 2 files changed, 25 insertions(+), 5 deletions(-)

diff --git a/nodes/nodes.py b/nodes/nodes.py
index cfec561..1472ec8 100644
--- a/nodes/nodes.py
+++ b/nodes/nodes.py
@@ -558,9 +558,9 @@ and performs comfy model management functions and garbage collection,
 reports free VRAM before and after the operations.
 """
 
-    def VRAMdebug(self, gc_collect,empty_cache, unload_all_models, image_pass=None, model_pass=None, any_input=None):
+    def VRAMdebug(self, gc_collect, empty_cache, unload_all_models, image_pass=None, model_pass=None, any_input=None):
         freemem_before = model_management.get_free_memory()
-        print("VRAMdebug: free memory before: ", freemem_before)
+        print("VRAMdebug: free memory before: ", f"{freemem_before:,.0f}")
         if empty_cache:
             model_management.soft_empty_cache()
         if unload_all_models:
@@ -569,9 +569,12 @@ reports free VRAM before and after the operations.
             import gc
             gc.collect()
         freemem_after = model_management.get_free_memory()
-        print("VRAMdebug: free memory after: ", freemem_after)
-        print("VRAMdebug: freed memory: ", freemem_after - freemem_before)
-        return (any_input, image_pass, model_pass, freemem_before, freemem_after)
+        print("VRAMdebug: free memory after: ", f"{freemem_after:,.0f}")
+        print("VRAMdebug: freed memory: ", f"{freemem_after - freemem_before:,.0f}")
+        return {"ui": {
+            "text": [f"{freemem_before:,.0f}x{freemem_after:,.0f}"]}, 
+            "result": (any_input, image_pass, model_pass, freemem_before, freemem_after) 
+        }
 
 class SomethingToString:
     @classmethod
diff --git a/web/js/jsnodes.js b/web/js/jsnodes.js
index c09e893..6131289 100644
--- a/web/js/jsnodes.js
+++ b/web/js/jsnodes.js
@@ -113,6 +113,23 @@ app.registerExtension({
 					return r
 				}
 				break;
+			case "VRAM_Debug":
+				const onVRAM_DebugConnectInput = nodeType.prototype.onConnectInput;
+				nodeType.prototype.onConnectInput = function (targetSlot, type, output, originNode, originSlot) {
+					const v = onVRAM_DebugConnectInput?.(this, arguments);
+					targetSlot.outputs[3]["name"] = "freemem_before"
+					targetSlot.outputs[4]["name"] = "freemem_after" 
+					return v;
+				}
+				const onVRAM_DebugExecuted = nodeType.prototype.onExecuted;
+				nodeType.prototype.onExecuted = function(message) {
+					const r = onVRAM_DebugExecuted? onVRAM_DebugExecuted.apply(this,arguments): undefined
+					let values = message["text"].toString().split('x');
+					this.outputs[3]["name"] = values[0] + "   freemem_before"
+					this.outputs[4]["name"] = values[1] + "      freemem_after" 
+					return r
+				}
+				break;
 
 			case "JoinStringMulti":
 				nodeType.prototype.onNodeCreated = function () {