Merge branch 'release/0.5.0'

Author: ty0x2333

CHANGELOG.md

@@ -1,3 +1,17 @@
+0.5.0
+---
+
+### Enhancement
+
+1. Add VRAM Used to the ExecutionTime Node.
+
+   PR: https://github.com/ty0x2333/ComfyUI-Dev-Utils/pull/20
+
+   Thanks [drake7707](https://github.com/drake7707) .
+
+2. Add a Max statistics row to the ExecutionTime Node to easily view the longest execution time of a node and the VRAM
+   requirements of the workflow.
+
 0.4.3
 ---
 

nodes/execution_time.py

@@ -1,7 +1,10 @@
 import time
 
+import torch
+
 import execution
 import server
+# import model_management
 
 
 class ExecutionTime:
@@ -22,20 +25,48 @@ def process(self):
 CURRENT_START_EXECUTION_DATA = None
 
 
+# def get_free_vram():
+#     dev = model_management.get_torch_device()
+#     if hasattr(dev, 'type') and (dev.type == 'cpu' or dev.type == 'mps'):
+#         return 0
+#     else:
+#         return model_management.get_free_memory(dev)
+
+
+def get_peak_memory():
+    if not torch.cuda.is_available():
+        return 0
+    device = torch.device('cuda')
+    return torch.cuda.max_memory_allocated(device)
+
+
+def reset_peak_memory_record():
+    if not torch.cuda.is_available():
+        return
+    device = torch.device('cuda')
+    torch.cuda.reset_max_memory_allocated(device)
+
+
 def handle_execute(class_type, last_node_id, prompt_id, server, unique_id):
     if not CURRENT_START_EXECUTION_DATA:
         return
     start_time = CURRENT_START_EXECUTION_DATA['nodes_start_perf_time'].get(unique_id)
+    start_vram = CURRENT_START_EXECUTION_DATA['nodes_start_vram'].get(unique_id)
     if start_time:
         end_time = time.perf_counter()
         execution_time = end_time - start_time
+
+        end_vram = get_peak_memory()
+        vram_used = end_vram - start_vram
+        print(f"end_vram - start_vram: {end_vram} - {start_vram} = {vram_used}")
         if server.client_id is not None and last_node_id != server.last_node_id:
             server.send_sync(
                 "TyDev-Utils.ExecutionTime.executed",
-                {"node": unique_id, "prompt_id": prompt_id, "execution_time": int(execution_time * 1000)},
+                {"node": unique_id, "prompt_id": prompt_id, "execution_time": int(execution_time * 1000),
+                 "vram_used": vram_used},
                 server.client_id
             )
-        print(f"#{unique_id} [{class_type}]: {execution_time:.2f}s")
+        print(f"#{unique_id} [{class_type}]: {execution_time:.2f}s - vram {vram_used}b")
 
 
 try:
@@ -91,7 +122,8 @@ def swizzle_send_sync(self, event, data, sid=None):
     if event == "execution_start":
         CURRENT_START_EXECUTION_DATA = dict(
             start_perf_time=time.perf_counter(),
-            nodes_start_perf_time={}
+            nodes_start_perf_time={},
+            nodes_start_vram={}
         )
 
     origin_func(self, event=event, data=data, sid=sid)
@@ -113,6 +145,8 @@ def swizzle_send_sync(self, event, data, sid=None):
         else:
             node_id = data.get("node")
             CURRENT_START_EXECUTION_DATA['nodes_start_perf_time'][node_id] = time.perf_counter()
+            reset_peak_memory_record()
+            CURRENT_START_EXECUTION_DATA['nodes_start_vram'][node_id] = get_peak_memory()
 
 
 server.PromptServer.send_sync = swizzle_send_sync

pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 name = "comfyui-dev-utils"
 description = "Execution Time Analysis, Reroute Enhancement, Node collection for developers."
-version = "0.4.3"
+version = "0.5.0"
 license = { text = "MIT License" }
 dependencies = ["aiohttp-sse"]
 

web/executionTime.js

@@ -27,6 +27,19 @@ function formatExecutionTime(time) {
     return `${(time / 1000.0).toFixed(2)}s`
 }
 
+// Reference: https://gist.github.com/zentala/1e6f72438796d74531803cc3833c039c
+function formatBytes(bytes, decimals) {
+    if (bytes === 0) {
+        return '0 B'
+    }
+    const k = 1024,
+        dm = decimals || 2,
+        sizes = ['Bytes', 'KB', 'MB', 'GB', 'TB', 'PB', 'EB', 'ZB', 'YB'],
+        i = Math.floor(Math.log(bytes) / Math.log(k));
+    return parseFloat((bytes / Math.pow(k, i)).toFixed(dm)) + ' ' + sizes[i];
+}
+
+
 // Reference: https://github.com/ltdrdata/ComfyUI-Manager/blob/main/js/comfyui-manager.js
 function drawBadge(node, orig, restArgs) {
     let ctx = restArgs[0];
@@ -35,7 +48,7 @@ function drawBadge(node, orig, restArgs) {
     if (!node.flags.collapsed && node.constructor.title_mode != LiteGraph.NO_TITLE) {
         let text = "";
         if (node.ty_et_execution_time !== undefined) {
-            text = formatExecutionTime(node.ty_et_execution_time);
+            text = formatExecutionTime(node.ty_et_execution_time) + " - vram " + formatBytes(node.ty_et_vram_used, 2);
         } else if (node.ty_et_start_time !== undefined) {
             text = formatExecutionTime(LiteGraph.getTime() - node.ty_et_start_time);
         }
@@ -140,7 +153,8 @@ function buildTableHtml() {
                 $el("th", {style: headerThStyle, "textContent": "Node Title"}),
                 $el("th", {style: headerThStyle, "textContent": "Current Time"}),
                 $el("th", {style: headerThStyle, "textContent": "Per Time"}),
-                $el("th", {style: headerThStyle, "textContent": "Current / Pre Diff"})
+                $el("th", {style: headerThStyle, "textContent": "Cur / Pre Time Diff"}),
+                $el("th", {style: headerThStyle, "textContent": "VRAM Used"})
             ])
         ]),
         tableBody,
@@ -170,6 +184,9 @@ function buildTableHtml() {
         return [diffColor, diffText]
     }
 
+    let max_execution_time = null
+    let max_vram_used = null
+
     runningData.nodes_execution_time.forEach(function (item) {
         const nodeId = item.node;
         const node = app.graph.getNodeById(nodeId)
@@ -178,7 +195,16 @@ function buildTableHtml() {
 
         const [diffColor, diffText] = diff(item.execution_time, preExecutionTime);
 
+        if (max_execution_time == null || item.execution_time > max_execution_time) {
+            max_execution_time = item.execution_time
+        }
+
+        if (max_vram_used == null || item.vram_used > max_vram_used) {
+            max_vram_used = item.vram_used
+        }
+
         tableBody.append($el("tr", {
+            style: {"cursor": "pointer"},
             onclick: () => {
                 if (node) {
                     app.canvas.selectNode(node, false);
@@ -199,10 +225,36 @@ function buildTableHtml() {
                 },
                 "textContent": diffText
             }),
+            $el("td", {style: {"textAlign": "right"}, "textContent": formatBytes(item.vram_used, 2)}),
         ]))
     });
     if (runningData.total_execution_time !== null) {
         const [diffColor, diffText] = diff(runningData.total_execution_time, lastRunningDate?.total_execution_time);
+
+         tableFooter.append($el("tr", [
+            $el("td", {style: {"textAlign": "right"}, "textContent": 'Max'}),
+            $el("td", {style: {"textAlign": "right"}, "textContent": ''}),
+            $el("td", {
+                style: {"textAlign": "right"},
+                "textContent": max_execution_time != null ? formatExecutionTime(max_execution_time) : ''
+            }),
+            $el("td", {
+                style: {"textAlign": "right"},
+                "textContent": ''
+            }),
+            $el("td", {
+                style: {
+                    "textAlign": "right",
+                    "color": diffColor
+                },
+                "textContent": ''
+            }),
+            $el("td", {
+                style: {"textAlign": "right"},
+                "textContent": max_vram_used != null ? formatBytes(max_vram_used, 2) : ''
+            }),
+        ]))
+
         tableFooter.append($el("tr", [
             $el("td", {style: {"textAlign": "right"}, "textContent": 'Total'}),
             $el("td", {style: {"textAlign": "right"}, "textContent": ''}),
@@ -221,6 +273,7 @@ function buildTableHtml() {
                 },
                 "textContent": diffText
             }),
+            $el("td", {style: {"textAlign": "right"}, "textContent": ""}),
         ]))
     }
     return table;
@@ -261,11 +314,13 @@ app.registerExtension({
             const node = app.graph.getNodeById(detail.node)
             if (node) {
                 node.ty_et_execution_time = detail.execution_time;
+                node.ty_et_vram_used = detail.vram_used;
             }
             const index = runningData.nodes_execution_time.findIndex(x => x.node === detail.node);
             const data = {
                 node: detail.node,
-                execution_time: detail.execution_time
+                execution_time: detail.execution_time,
+                vram_used: detail.vram_used
             };
             if (index > 0) {
                 runningData.nodes_execution_time[index] = data
@@ -283,6 +338,7 @@ app.registerExtension({
             app.graph._nodes.forEach(function (node) {
                 delete node.ty_et_start_time
                 delete node.ty_et_execution_time
+                delete node.ty_et_vram_used
             });
             runningData = {
                 nodes_execution_time: [],
@@ -337,7 +393,7 @@ app.registerExtension({
                 const thUnscaledHeight = 24;
                 const tableUnscaledHeight = thUnscaledHeight * tableHeight / thHeight;
                 const autoResizeMaxHeight = 300;
-                return [Math.max(originSize[0], 480), originSize[1] + Math.min(tableUnscaledHeight, autoResizeMaxHeight) - LiteGraph.NODE_WIDGET_HEIGHT];
+                return [Math.max(originSize[0], 600), originSize[1] + Math.min(tableUnscaledHeight, autoResizeMaxHeight) - LiteGraph.NODE_WIDGET_HEIGHT];
             }
 
             const nodeCreated = nodeType.prototype.onNodeCreated;