fix: upload status issues, and gif uploads

This commit is contained in:
BennyKok 2024-01-12 14:53:42 +08:00
parent 547d645a0c
commit ac6e863b95
3 changed files with 50 additions and 37 deletions

View File

@ -79,7 +79,7 @@ COMFY_API_AVAILABLE_MAX_RETRIES = 500
# Time to wait between poll attempts in milliseconds
COMFY_POLLING_INTERVAL_MS = 250
# Maximum number of poll attempts
COMFY_POLLING_MAX_RETRIES = 500
COMFY_POLLING_MAX_RETRIES = 1000
# Host where ComfyUI is running
COMFY_HOST = "127.0.0.1:8188"

View File

@ -200,12 +200,10 @@ prompt_server = server.PromptServer.instance
send_json = prompt_server.send_json
async def send_json_override(self, event, data, sid=None):
# print("INTERNAL:", event, data, sid)
prompt_id = data.get('prompt_id')
# now we send everything
await send(event, data)
await self.send_json_original(event, data, sid)
await asyncio.wait([send(event, data), self.send_json_original(event, data, sid)])
if event == 'execution_start':
update_run(prompt_id, Status.RUNNING)
@ -217,12 +215,13 @@ async def send_json_override(self, event, data, sid=None):
update_run(prompt_id, Status.SUCCESS)
if event == 'execution_error':
# Careful this might not be fully awaited.
await update_run_with_output(prompt_id, data)
update_run(prompt_id, Status.FAILED)
asyncio.create_task(update_run_with_output(prompt_id, data))
# await update_run_with_output(prompt_id, data)
if event == 'executed' and 'node' in data and 'output' in data:
asyncio.create_task(update_run_with_output(prompt_id, data.get('output'), node_id=data.get('node')))
await update_run_with_output(prompt_id, data.get('output'), node_id=data.get('node'))
# await update_run_with_output(prompt_id, data.get('output'), node_id=data.get('node'))
# update_run_with_output(prompt_id, data.get('output'))
@ -329,6 +328,25 @@ def is_prompt_done(prompt_id):
return False
# Use to handle upload error and send back to ComfyDeploy
async def handle_error(prompt_id, data, e: Exception):
error_type = type(e).__name__
stack_trace = traceback.format_exc().strip()
body = {
"run_id": prompt_id,
"output_data": {
"error": {
"type": error_type,
"message": str(e),
"stack_trace": stack_trace
}
}
}
await update_file_status(prompt_id, data, False, have_error=True)
print(body)
print(f"Error occurred while uploading file: {e}")
# Mark the current prompt requires upload, and block it from being marked as success
async def update_file_status(prompt_id, data, uploading, have_error=False, node_id=None):
if 'uploading_nodes' not in prompt_metadata[prompt_id]:
prompt_metadata[prompt_id]['uploading_nodes'] = set()
@ -366,6 +384,26 @@ async def update_file_status(prompt_id, data, uploading, have_error=False, node_
"prompt_id": prompt_id,
})
# Upload files in the background
async def upload_in_background(prompt_id, data, node_id=None, have_upload=True):
try:
images = data.get('images', [])
for image in images:
await upload_file(prompt_id, image.get("filename"), subfolder=image.get("subfolder"), type=image.get("type"), content_type=image.get("content_type", "image/png"))
files = data.get('files', [])
for file in files:
await upload_file(prompt_id, file.get("filename"), subfolder=file.get("subfolder"), type=file.get("type"), content_type=file.get("content_type", "image/png"))
gifs = data.get('gifs', [])
for gif in gifs:
await upload_file(prompt_id, gif.get("filename"), subfolder=gif.get("subfolder"), type=gif.get("type"), content_type=gif.get("format", "image/gif"))
if have_upload:
await update_file_status(prompt_id, data, False, node_id=node_id)
except Exception as e:
await handle_error(prompt_id, data, e)
async def update_run_with_output(prompt_id, data, node_id=None):
if prompt_id in prompt_metadata:
status_endpoint = prompt_metadata[prompt_id]['status_endpoint']
@ -377,43 +415,16 @@ async def update_run_with_output(prompt_id, data, node_id=None):
try:
have_upload = 'images' in data or 'files' in data or 'gifs' in data
print("\nhave_upload", have_upload, node_id)
if have_upload:
await update_file_status(prompt_id, data, True, node_id=node_id)
images = data.get('images', [])
for image in images:
await upload_file(prompt_id, image.get("filename"), subfolder=image.get("subfolder"), type=image.get("type"), content_type=image.get("content_type", "image/png"))
files = data.get('files', [])
for file in files:
await upload_file(prompt_id, file.get("filename"), subfolder=file.get("subfolder"), type=file.get("type"), content_type=file.get("content_type", "image/png"))
gifs = data.get('gifs', [])
for gif in gifs:
await upload_file(prompt_id, file.get("filename"), subfolder=file.get("subfolder"), type=file.get("type"), content_type=file.get("format", "image/gif"))
if have_upload:
await update_file_status(prompt_id, data, False, node_id=node_id)
asyncio.create_task(upload_in_background(prompt_id, data, node_id=node_id, have_upload=have_upload))
except Exception as e:
error_type = type(e).__name__
stack_trace = traceback.format_exc().strip()
body = {
"run_id": prompt_id,
"output_data": {
"error": {
"type": error_type,
"message": str(e),
"stack_trace": stack_trace
}
}
}
await update_file_status(prompt_id, data, False, have_error=True)
print(body)
print(f"Error occurred while uploading file: {e}")
await handle_error(prompt_id, data, e)
requests.post(status_endpoint, json=body)

View File

@ -23,7 +23,9 @@ export async function RunOutputs({ run_id }: { run_id: string }) {
<TableBody>
{outputs?.map((run) => {
const fileName =
run.data.images?.[0].filename || run.data.files?.[0].filename;
run.data.images?.[0].filename ||
run.data.files?.[0].filename ||
run.data.gifs?.[0].filename;
if (!fileName)
return (