From 97096a90351eb14a2bf82ad0ffd18e71b810ea9e Mon Sep 17 00:00:00 2001
From: bennykok <itechbenny@gmail.com>
Date: Sun, 25 Feb 2024 22:47:54 -0800
Subject: [PATCH] feat(plugin): send live_status and elapsed_time

---
 custom_routes.py | 18 ++++++++++++++++++
 1 file changed, 18 insertions(+)

diff --git a/custom_routes.py b/custom_routes.py
index 8de4d37..7818280 100644
--- a/custom_routes.py
+++ b/custom_routes.py
@@ -55,6 +55,7 @@ class SimplePrompt(BaseModel):
     uploading_nodes: set = set()
     done: bool = False
     is_realtime: bool = False,
+    start_time: Optional[float] = None,
     
 streaming_prompt_metadata: dict[str, StreamingPrompt] = {}
 
@@ -499,12 +500,24 @@ async def send_json_override(self, event, data, sid=None):
 
     if event == 'execution_start':
         update_run(prompt_id, Status.RUNNING)
+        
+        if prompt_id in prompt_metadata:
+            prompt_metadata[prompt_id].start_time = time.perf_counter()
 
     # 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)
         if not have_pending_upload(prompt_id):
             update_run(prompt_id, Status.SUCCESS)
+            if prompt_id in prompt_metadata:
+                current_time = time.perf_counter()
+                if prompt_metadata[prompt_id].start_time is not None:
+                    elapsed_time = current_time - prompt_metadata[prompt_id].start_time
+                    print(f"Elapsed time: {elapsed_time} seconds")
+                    await 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')
@@ -522,6 +535,11 @@ 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']
             print("updating run live status", class_type)
+            await send("live_status", {
+                "prompt_id": prompt_id,
+                "current_node": class_type,
+                "progress": calculated_progress,
+            }, sid=sid)
             await update_run_live_status(prompt_id, "Executing " + class_type, calculated_progress)
 
     if event == 'execution_cached' and data.get('nodes') is not None: