fix: simply js import

I tracked the time and takes ~200ms everytime that we send the "Executing <NODE NAME> n%".
So this means that if you have 10 custom nodes we are adding 2 extra seconds to the execution.
200 * 10 = 2,000.
Some workflows are more complext and have more custom nodes, so this only keeps increasing.

custom_routes.py

@@ -26,7 +26,10 @@ import copy
 import struct
 from aiohttp import web, ClientSession, ClientError, ClientTimeout, ClientResponseError
 import atexit
-
+from model_management import get_torch_device
+import torch
+import psutil
+from collections import OrderedDict
 # Global session
 client_session = None
 
@@ -1120,9 +1123,138 @@ async def proxy_to_comfydeploy(request):
 
 
 prompt_server = server.PromptServer.instance
-send_json = prompt_server.send_json
 
 
+NODE_EXECUTION_TIMES = {} # New dictionary to store node execution times
+CURRENT_START_EXECUTION_DATA = None
+
+def get_peak_memory():
+    device = get_torch_device()
+    if device.type == 'cuda':
+        return torch.cuda.max_memory_allocated(device)
+    elif device.type == 'mps':
+        # Return system memory usage for MPS devices
+        return psutil.Process().memory_info().rss
+    return 0
+
+
+def reset_peak_memory_record():
+    device = get_torch_device()
+    if device.type == 'cuda':
+        torch.cuda.reset_max_memory_allocated(device)
+    # MPS doesn't need reset as we're not tracking its memory
+
+
+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
+        global NODE_EXECUTION_TIMES
+        # print(f"end_vram - start_vram: {end_vram} - {start_vram} = {vram_used}")
+        NODE_EXECUTION_TIMES[unique_id] = {
+            "time": execution_time,
+            "class_type": class_type,
+            "vram_used": vram_used
+        }
+        # print(f"#{unique_id} [{class_type}]: {execution_time:.2f}s - vram {vram_used}b")
+
+
+try:
+    origin_execute = execution.execute
+
+    def swizzle_execute(
+        server,
+        dynprompt,
+        caches,
+        current_item,
+        extra_data,
+        executed,
+        prompt_id,
+        execution_list,
+        pending_subgraph_results,
+    ):
+        unique_id = current_item
+        class_type = dynprompt.get_node(unique_id)["class_type"]
+        last_node_id = server.last_node_id
+        result = origin_execute(
+            server,
+            dynprompt,
+            caches,
+            current_item,
+            extra_data,
+            executed,
+            prompt_id,
+            execution_list,
+            pending_subgraph_results,
+        )
+        handle_execute(class_type, last_node_id, prompt_id, server, unique_id)
+        return result
+
+    execution.execute = swizzle_execute
+except Exception as e:
+    pass
+
+def format_table(headers, data):
+    # Calculate column widths
+    widths = [len(h) for h in headers]
+    for row in data:
+        for i, cell in enumerate(row):
+            widths[i] = max(widths[i], len(str(cell)))
+    
+    # Create separator line
+    separator = '+' + '+'.join('-' * (w + 2) for w in widths) + '+'
+    
+    # Format header
+    result = [separator]
+    header_row = '|' + '|'.join(f' {h:<{w}} ' for w, h in zip(widths, headers)) + '|'
+    result.append(header_row)
+    result.append(separator)
+    
+    # Format data rows
+    for row in data:
+        data_row = '|' + '|'.join(f' {str(cell):<{w}} ' for w, cell in zip(widths, row)) + '|'
+        result.append(data_row)
+    
+    result.append(separator)
+    return '\n'.join(result)
+
+origin_func = server.PromptServer.send_sync
+def swizzle_send_sync(self, event, data, sid=None):
+    # print(f"swizzle_send_sync, event: {event}, data: {data}")
+    global CURRENT_START_EXECUTION_DATA
+    if event == "execution_start":
+        global NODE_EXECUTION_TIMES
+        NODE_EXECUTION_TIMES = {}  # Reset execution times at start
+        CURRENT_START_EXECUTION_DATA = dict(
+            start_perf_time=time.perf_counter(),
+            nodes_start_perf_time={},
+            nodes_start_vram={},
+        )
+
+    origin_func(self, event=event, data=data, sid=sid)
+
+    if event == "executing" and data and CURRENT_START_EXECUTION_DATA:
+        if data.get("node") is not None:
+            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
+
+send_json = prompt_server.send_json
+
 async def send_json_override(self, event, data, sid=None):
     # logger.info("INTERNAL:", event, data, sid)
     prompt_id = data.get("prompt_id")
@@ -1145,11 +1277,61 @@ async def send_json_override(self, event, data, sid=None):
     asyncio.create_task(update_run_ws_event(prompt_id, event, data))
 
     if event == "execution_start":
-        await update_run(prompt_id, Status.RUNNING)
-
         if prompt_id in prompt_metadata:
             prompt_metadata[prompt_id].start_time = time.perf_counter()
             
+        asyncio.create_task(update_run(prompt_id, Status.RUNNING))
+        
+        
+    if event == "executing" and data and CURRENT_START_EXECUTION_DATA:
+        if data.get("node") is None:
+            start_perf_time = CURRENT_START_EXECUTION_DATA.get("start_perf_time")
+            new_data = data.copy()
+            if start_perf_time is not None:
+                execution_time = time.perf_counter() - start_perf_time
+                new_data["execution_time"] = int(execution_time * 1000)
+                
+            # Replace the print statements with tabulate
+            headers = ["Node ID", "Type", "Time (s)", "VRAM (GB)"]
+            table_data = []
+            node_execution_array = []  # New array to store execution data
+            
+            for node_id, node_data in NODE_EXECUTION_TIMES.items():
+                vram_gb = node_data['vram_used'] / (1024**3)  # Convert bytes to GB
+                table_data.append([
+                    f"#{node_id}",
+                    node_data['class_type'],
+                    f"{node_data['time']:.2f}",
+                    f"{vram_gb:.2f}"
+                ])
+                
+                # Add to our new array format
+                node_execution_array.append({
+                    "id": node_id,
+                    **node_data,
+                })
+            
+            # Add total execution time as the last row
+            table_data.append([
+                "TOTAL",
+                "-",
+                f"{execution_time:.2f}",
+                "-"
+            ])
+            
+            prompt_id = data.get("prompt_id")
+            asyncio.create_task(update_run_with_output(
+                prompt_id,
+                node_execution_array,  # Send the array instead of the OrderedDict
+            ))
+            
+            print(node_execution_array)
+            
+            # print("\n=== Node Execution Times ===")
+            logger.info("Printing Node Execution Times")
+            logger.info(format_table(headers, table_data))
+            # print("========================\n")
+
     # the last executing event is none, then the workflow is finished
     if event == "executing" and data.get("node") is None:
         mark_prompt_done(prompt_id=prompt_id)
@@ -1160,11 +1342,11 @@ async def send_json_override(self, event, data, sid=None):
                 if prompt_metadata[prompt_id].start_time is not None:
                     elapsed_time = current_time - prompt_metadata[prompt_id].start_time
                     logger.info(f"Elapsed time: {elapsed_time} seconds")
-                    await send(
+                    asyncio.create_task(send(
                         "elapsed_time",
                         {"prompt_id": prompt_id, "elapsed_time": elapsed_time},
                         sid=sid,
-                    )
+                    ))
 
     if event == "executing" and data.get("node") is not None:
         node = data.get("node")
@@ -1188,7 +1370,7 @@ async def send_json_override(self, event, data, sid=None):
             prompt_metadata[prompt_id].last_updated_node = node
             class_type = prompt_metadata[prompt_id].workflow_api[node]["class_type"]
             logger.info(f"At: {round(calculated_progress * 100)}% - {class_type}")
-            await send(
+            asyncio.create_task(send(
                 "live_status",
                 {
                     "prompt_id": prompt_id,
@@ -1196,10 +1378,10 @@ async def send_json_override(self, event, data, sid=None):
                     "progress": calculated_progress,
                 },
                 sid=sid,
-            )
-            await update_run_live_status(
+            ))
+            asyncio.create_task(update_run_live_status(
                 prompt_id, "Executing " + class_type, calculated_progress
-            )
+            ))
 
     if event == "execution_cached" and data.get("nodes") is not None:
         if prompt_id in prompt_metadata:
@@ -1227,7 +1409,8 @@ async def send_json_override(self, event, data, sid=None):
                 "node_class": class_type,
             }
             if class_type == "PreviewImage":
-                logger.info("Skipping preview image")
+                pass
+                # logger.info("Skipping preview image")
             else:
                 await update_run_with_output(
                     prompt_id,
@@ -1239,9 +1422,10 @@ async def send_json_override(self, event, data, sid=None):
                     comfy_message_queues[prompt_id].put_nowait(
                         {"event": "output_ready", "data": data}
                     )
-            logger.info(f"Executed {class_type} {data}")
+            # logger.info(f"Executed {class_type} {data}")
         else:
-            logger.info(f"Executed {data}")
+            pass
+            # logger.info(f"Executed {data}")
 
 
 # Global variable to keep track of the last read line number

requirements.txt

@@ -3,4 +3,5 @@ pydantic
 opencv-python
 imageio-ffmpeg
 brotli
+tabulate
 # logfire

web-plugin/api.js

@@ -1,4 +0,0 @@
-/** @typedef {import('../../../web/scripts/api.js').api} API*/
-import { api as _api } from '../../scripts/api.js';
-/** @type {API} */
-export const api = _api;

web-plugin/app.js

@@ -1,4 +0,0 @@
-/** @typedef {import('../../../web/scripts/app.js').ComfyApp} ComfyApp*/
-import { app as _app } from '../../scripts/app.js';
-/** @type {ComfyApp} */
-export const app = _app;

web-plugin/index.js

@@ -1,8 +1,11 @@
-import { app } from "./app.js";
-import { api } from "./api.js";
-import { ComfyWidgets, LGraphNode } from "./widgets.js";
+import { app } from "../../scripts/app.js";
+import { api } from "../../scripts/api.js";
+// import { LGraphNode } from "../../scripts/widgets.js";
+LGraphNode = LiteGraph.LGraphNode;
+import { ComfyDialog, $el } from "../../scripts/ui.js";
+
 import { generateDependencyGraph } from "https://esm.sh/[email protected]";
-import { ComfyDeploy } from "https://esm.sh/comfydeploy@0.0.19-beta.30";
+import { ComfyDeploy } from "https://esm.sh/comfydeploy@2.0.0-beta.69";
 
 const styles = `
 .comfydeploy-menu-item {
@@ -1287,8 +1290,6 @@ function addButton() {
 
 app.registerExtension(ext);
 
-import { ComfyDialog, $el } from "../../scripts/ui.js";
-
 export class InfoDialog extends ComfyDialog {
   constructor() {
     super();

web-plugin/widgets.js

@@ -1,18 +0,0 @@
-// /** @typedef {import('../../../web/scripts/api.js').api} API*/
-// import { api as _api } from "../../scripts/api.js";
-// /** @type {API} */
-// export const api = _api;
-
-/** @typedef {typeof import('../../../web/scripts/widgets.js').ComfyWidgets} Widgets*/
-import { ComfyWidgets as _ComfyWidgets } from "../../scripts/widgets.js";
-
-/**
- * @type {Widgets}
- */
-export const ComfyWidgets = _ComfyWidgets;
-
-// import { LGraphNode as _LGraphNode } from "../../types/litegraph.js";
-
-/** @typedef {typeof import('../../../web/types/litegraph.js').LGraphNode} LGraphNode*/
-/** @type {LGraphNode}*/
-export const LGraphNode = LiteGraph.LGraphNode;