chunzhi/comfyui-deploy
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: send ws after cd logic

Browse Source
This commit is contained in:
bennykok 2024-10-04 16:03:53 -07:00
parent 2f33bcf497
commit c115c22a91
1 changed files with 16 additions and 23 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

39
custom_routes.py
View File

@ -395,7 +395,7 @@ async def comfy_deploy_run(request):
# We proxy the request to Comfy Deploy, this is a native run # We proxy the request to Comfy Deploy, this is a native run
if "is_native_run" in data: if "is_native_run" in data:
async with aiohttp.ClientSession() as session: async with aiohttp.ClientSession() as session:
pprint(data) # pprint(data)
# headers = request.headers.copy() # headers = request.headers.copy()
# headers['Content-Type'] = 'application/json' # headers['Content-Type'] = 'application/json'
async with session.post(data.get("native_run_api_endpoint"), json=data, headers={ async with session.post(data.get("native_run_api_endpoint"), json=data, headers={
@ -403,7 +403,7 @@ async def comfy_deploy_run(request):
'Authorization': request.headers.get('Authorization') 'Authorization': request.headers.get('Authorization')
}) as response: }) as response:
data = await response.json() data = await response.json()
print(data) # print(data)
if "cd_token" in data: if "cd_token" in data:
token = data["cd_token"] token = data["cd_token"]
@ -1002,24 +1002,14 @@ async def send_json_override(self, event, data, sid=None):
target_sid = sid target_sid = sid
if target_sid == "comfy_deploy_instance": if target_sid == "comfy_deploy_instance":
target_sid = None target_sid = None
# now we send everything # now we send everything
await asyncio.wait([ await asyncio.wait([
asyncio.create_task(send(event, data, sid=target_sid)), asyncio.create_task(send(event, data, sid=target_sid)),
asyncio.create_task(self.send_json_original(event, data, sid)) asyncio.create_task(self.send_json_original(event, data, sid))
]) ])
if prompt_id in comfy_message_queues:
comfy_message_queues[prompt_id].put_nowait({
"event": event,
"data": data
})
asyncio.create_task(update_run_ws_event(prompt_id, event, data)) asyncio.create_task(update_run_ws_event(prompt_id, event, data))
# event_emitter.emit("send_json", {
# "event": event,
# "data": data
# })
if event == 'execution_start': if event == 'execution_start':
await update_run(prompt_id, Status.RUNNING) await update_run(prompt_id, Status.RUNNING)
@ -1058,7 +1048,7 @@ async def send_json_override(self, event, data, sid=None):
return return
prompt_metadata[prompt_id].last_updated_node = node prompt_metadata[prompt_id].last_updated_node = node
class_type = prompt_metadata[prompt_id].workflow_api[node]['class_type'] class_type = prompt_metadata[prompt_id].workflow_api[node]['class_type']
logger.info(f"At: {calculated_progress * 100}% - {class_type}") logger.info(f"At: {round(calculated_progress * 100)}% - {class_type}")
await send("live_status", { await send("live_status", {
"prompt_id": prompt_id, "prompt_id": prompt_id,
"current_node": class_type, "current_node": class_type,
@ -1087,7 +1077,6 @@ async def send_json_override(self, event, data, sid=None):
if prompt_id in prompt_metadata: if prompt_id in prompt_metadata:
node = data.get('node') node = data.get('node')
class_type = prompt_metadata[prompt_id].workflow_api[node]['class_type'] class_type = prompt_metadata[prompt_id].workflow_api[node]['class_type']
logger.info(f"Executed {class_type} {data}")
node_meta = { node_meta = {
"node_id": node, "node_id": node,
"node_class": class_type, "node_class": class_type,
@ -1095,12 +1084,16 @@ async def send_json_override(self, event, data, sid=None):
if class_type == "PreviewImage": if class_type == "PreviewImage":
logger.info("Skipping preview image") logger.info("Skipping preview image")
return return
await update_run_with_output(prompt_id, data.get('output'), node_id=data.get('node'), node_meta=node_meta)
logger.info(f"Executed {class_type} {data}")
else: else:
logger.info(f"Executed {data}") logger.info(f"Executed {data}")
await update_run_with_output(prompt_id, data.get('output'), node_id=data.get('node'), node_meta=node_meta) if prompt_id in comfy_message_queues:
# await update_run_with_output(prompt_id, data.get('output'), node_id=data.get('node')) comfy_message_queues[prompt_id].put_nowait({
# update_run_with_output(prompt_id, data.get('output')) "event": event,
"data": data
})
# Global variable to keep track of the last read line number # Global variable to keep track of the last read line number
last_read_line_number = 0 last_read_line_number = 0
@ -1262,7 +1255,7 @@ async def upload_with_retry(session, url, headers, data, max_retries=3, initial_
response.raise_for_status() # This will raise an exception for 4xx and 5xx status codes response.raise_for_status() # This will raise an exception for 4xx and 5xx status codes
response_text = await response.text() response_text = await response.text()
logger.info(f"Response body: {response_text[:1000]}...") # logger.info(f"Response body: {response_text[:1000]}...")
logger.info("Upload successful") logger.info("Upload successful")
return response # Successful upload, exit the retry loop return response # Successful upload, exit the retry loop
@ -1319,13 +1312,13 @@ async def upload_file(prompt_id, filename, subfolder=None, content_type="image/p
target_url = f"{file_upload_endpoint}?file_name={filename}&run_id={prompt_id}&type={content_type}&version=v2" target_url = f"{file_upload_endpoint}?file_name={filename}&run_id={prompt_id}&type={content_type}&version=v2"
start_time = time.time() # Start timing here start_time = time.time() # Start timing here
logger.info(f"Target URL: {target_url}") # logger.info(f"Target URL: {target_url}")
result = await async_request_with_retry("GET", target_url, disable_timeout=True, token=token) result = await async_request_with_retry("GET", target_url, disable_timeout=True, token=token)
end_time = time.time() # End timing after the request is complete end_time = time.time() # End timing after the request is complete
logger.info("Time taken for getting file upload endpoint: {:.2f} seconds".format(end_time - start_time)) logger.info("Time taken for getting file upload endpoint: {:.2f} seconds".format(end_time - start_time))
ok = await result.json() ok = await result.json()
logger.info(f"Result: {ok}") # logger.info(f"Result: {ok}")
async with aiofiles.open(file, 'rb') as f: async with aiofiles.open(file, 'rb') as f:
data = await f.read() data = await f.read()
@ -1338,7 +1331,7 @@ async def upload_file(prompt_id, filename, subfolder=None, content_type="image/p
"Content-Length": size, "Content-Length": size,
} }
logger.info(headers) # logger.info(headers)
if ok.get('include_acl') is True: if ok.get('include_acl') is True:
headers["x-amz-acl"] = "public-read" headers["x-amz-acl"] = "public-read"
@ -1542,7 +1535,7 @@ async def update_run_with_output(prompt_id, data, node_id=None, node_meta=None):
"output_data": data, "output_data": data,
"node_meta": node_meta, "node_meta": node_meta,
} }
pprint(body) # pprint(body)
have_upload_media = False have_upload_media = False
if data is not None: if data is not None:
have_upload_media = 'images' in data or 'files' in data or 'gifs' in data or 'mesh' in data have_upload_media = 'images' in data or 'files' in data or 'gifs' in data or 'mesh' in data