From 0f58fbcebd4dc3a8903354bd62322335daf3ef2f Mon Sep 17 00:00:00 2001 From: bennykok Date: Mon, 5 Feb 2024 23:28:09 +0800 Subject: [PATCH] chore: update add live status --- custom_routes.py | 24 ++++++++++++++++++++++++ 1 file changed, 24 insertions(+) diff --git a/custom_routes.py b/custom_routes.py index 4416533..759a41f 100644 --- a/custom_routes.py +++ b/custom_routes.py @@ -101,6 +101,7 @@ async def comfy_deploy_run(request): prompt_metadata[prompt_id] = { 'status_endpoint': data.get('status_endpoint'), 'file_upload_endpoint': data.get('file_upload_endpoint'), + 'workflow_api': workflow_api } try: @@ -355,6 +356,15 @@ async def send_json_override(self, event, data, sid=None): if not have_pending_upload(prompt_id): update_run(prompt_id, Status.SUCCESS) + if event == 'executing' and data.get('node') is not None: + node = data.get('node') + if 'last_updated_node' in prompt_metadata[prompt_id] and prompt_metadata[prompt_id]['last_updated_node'] == node: + return + 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 update_run_live_status(prompt_id, "Executing " + class_type) + if event == 'execution_error': # Careful this might not be fully awaited. await update_run_with_output(prompt_id, data) @@ -377,6 +387,20 @@ class Status(Enum): # Global variable to keep track of the last read line number last_read_line_number = 0 +async def update_run_live_status(prompt_id, live_status): + if prompt_id not in prompt_metadata: + return + + status_endpoint = prompt_metadata[prompt_id]['status_endpoint'] + body = { + "run_id": prompt_id, + "live_status": live_status, + } + # requests.post(status_endpoint, json=body) + async with aiohttp.ClientSession() as session: + await session.post(status_endpoint, json=body) + + def update_run(prompt_id, status: Status): global last_read_line_number