diff --git a/nodes/execution_time.py b/nodes/execution_time.py index ffc5d68..c39b610 100644 --- a/nodes/execution_time.py +++ b/nodes/execution_time.py @@ -2,7 +2,7 @@ import time import execution import server - +import model_management class ExecutionTime: CATEGORY = "TyDev-Utils/Debug" @@ -26,16 +26,21 @@ 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_time'].get(unique_id) if start_time: end_time = time.perf_counter() + end_vram = model_management.get_free_memory() execution_time = end_time - start_time + # use abs because it can be negative, that means the model_manager already loaded the model when we started profiling and cleared the model at the end + # it's not super accurate for anything small but for large models it's really handy to know what the footprint is + vram_used = abs(end_vram - start_vram) 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 +96,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_time={} ) origin_func(self, event=event, data=data, sid=sid) @@ -113,6 +119,7 @@ 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() + CURRENT_START_EXECUTION_DATA['nodes_start_vram_time'][node_id] = model_management.get_free_memory() server.PromptServer.send_sync = swizzle_send_sync diff --git a/web/executionTime.js b/web/executionTime.js index 31a5841..ff5cc04 100644 --- a/web/executionTime.js +++ b/web/executionTime.js @@ -27,6 +27,17 @@ 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'; + var 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 +46,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 +151,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": "Current / Pre Diff"}), + $el("th", {style: headerThStyle, "textContent": "VRAM used"}) ]) ]), tableBody, @@ -199,6 +211,7 @@ function buildTableHtml() { }, "textContent": diffText }), + $el("td", {style: {"textAlign": "right"}, "textContent": formatBytes(item.vram_used,2)}), ])) }); if (runningData.total_execution_time !== null) { @@ -221,6 +234,7 @@ function buildTableHtml() { }, "textContent": diffText }), + $el("td", {style: {"textAlign": "right"}, "textContent": ""}), ])) } return table; @@ -261,11 +275,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 +299,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: [],