chore: update add live status

This commit is contained in:
bennykok 2024-02-05 23:28:09 +08:00
parent 6dc964c425
commit 0f58fbcebd

View File

@ -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