Add ComfyUI node timing for full workflow

app.py

@@ -16,6 +16,7 @@ import gradio as gr
 import requests
 import spaces
 import torch
+import websocket
 
 from scripts.workflow_client import load_workflow, patch_voicegate_workflow
 
@@ -153,10 +154,10 @@ def write_sine_wav(filename: str, *, seconds: float = 1.0, frequency: float = 44
     return filename
 
 
-def submit_prompt(workflow: dict[str, Any]) -> str:
+def submit_prompt(workflow: dict[str, Any], *, client_id: str | None = None) -> str:
     response = requests.post(
         f"{COMFY_URL}/prompt",
-        json={"prompt": workflow, "client_id": str(uuid.uuid4())},
+        json={"prompt": workflow, "client_id": client_id or str(uuid.uuid4())},
         timeout=120,
     )
     if not response.ok:
@@ -164,6 +165,83 @@ def submit_prompt(workflow: dict[str, Any]) -> str:
     return response.json()["prompt_id"]
 
 
+def execute_prompt_with_timing(workflow: dict[str, Any], *, timeout: float) -> tuple[str, dict[str, Any], list[str]]:
+    client_id = str(uuid.uuid4())
+    websocket_url = f"ws://{COMFY_HOST}:{COMFY_PORT}/ws?clientId={client_id}"
+    ws = websocket.create_connection(websocket_url, timeout=30)
+    prompt_id = submit_prompt(workflow, client_id=client_id)
+    started = time.time()
+    deadline = started + timeout
+    current_node: str | None = None
+    current_started = 0.0
+    node_durations: dict[str, float] = {}
+    node_order: list[str] = []
+    event_lines = [f"prompt_id={prompt_id}", "node_timing=started"]
+
+    def close_current_node(now: float) -> None:
+        nonlocal current_node, current_started
+        if current_node is not None:
+            node_durations[current_node] = node_durations.get(current_node, 0.0) + max(0.0, now - current_started)
+            current_node = None
+            current_started = 0.0
+
+    try:
+        while time.time() < deadline:
+            ws.settimeout(max(1.0, min(10.0, deadline - time.time())))
+            try:
+                message = ws.recv()
+            except websocket.WebSocketTimeoutException:
+                continue
+            if isinstance(message, bytes):
+                message = message.decode("utf-8", errors="replace")
+            try:
+                payload = json.loads(message)
+            except json.JSONDecodeError:
+                continue
+            event_type = payload.get("type")
+            data = payload.get("data") or {}
+            if data.get("prompt_id") not in (None, prompt_id):
+                continue
+
+            now = time.time()
+            if event_type == "executing":
+                close_current_node(now)
+                node = data.get("node")
+                if node is None:
+                    continue
+                current_node = str(node)
+                current_started = now
+                if current_node not in node_order:
+                    node_order.append(current_node)
+            elif event_type == "execution_success":
+                close_current_node(now)
+                event_lines.append(f"websocket_elapsed_sec={now - started:.1f}")
+                break
+            elif event_type == "execution_error":
+                close_current_node(now)
+                event_lines.append("websocket_execution_error:")
+                event_lines.append(json.dumps(data, ensure_ascii=False, indent=2)[:4000])
+                break
+        else:
+            close_current_node(time.time())
+            raise TimeoutError(f"Timed out waiting for prompt {prompt_id}")
+    finally:
+        ws.close()
+
+    history = wait_for_history(prompt_id, timeout=30)
+    timed_nodes = sorted(
+        ((node_id, node_durations.get(node_id, 0.0)) for node_id in node_order),
+        key=lambda item: item[1],
+        reverse=True,
+    )
+    if timed_nodes:
+        event_lines.append("node_timing_top:")
+        for node_id, seconds in timed_nodes[:20]:
+            class_type = workflow.get(node_id, {}).get("class_type", "unknown")
+            event_lines.append(f"{node_id} {class_type}: {seconds:.1f}s")
+    return prompt_id, history, event_lines
+
+
 def wait_for_history(prompt_id: str, timeout: float = 1200) -> dict[str, Any]:
     deadline = time.time() + timeout
     while time.time() < deadline:
@@ -492,7 +570,7 @@ def asr_gpu_test(audio_path: str | None) -> str:
     return "\n".join(lines)
 
 
-@spaces.GPU(duration=800)
+@spaces.GPU(duration=900)
 def full_voicegate_gpu_test(audio_path: str | None, target_language: str) -> str:
     lines = gpu_status_lines()
     started = time.time()
@@ -507,9 +585,8 @@ def full_voicegate_gpu_test(audio_path: str | None, target_language: str) -> str
         lines.append(f"input_audio={audio_filename}")
         lines.append(f"target_language={target_language}")
         prompt = full_voicegate_workflow(audio_filename, prefix, target_language or "English")
-        prompt_id = submit_prompt(prompt)
-        lines.append(f"prompt_id={prompt_id}")
-        history = wait_for_history(prompt_id, timeout=800)
+        _prompt_id, history, timing_lines = execute_prompt_with_timing(prompt, timeout=880)
+        lines.extend(timing_lines)
         lines.extend(history_summary(history))
         lines.append(f"elapsed_sec={time.time() - started:.1f}")
     except Exception as exc: