feat: experiment with await + asyncio.gather for multi file in same node

This commit is contained in:
bennykok 2024-09-17 12:56:42 -07:00
parent 946571e32e
commit c08b68c41f

View File

@ -1359,7 +1359,7 @@ async def update_file_status(prompt_id: str, data, uploading, have_error=False,
async def handle_upload(prompt_id: str, data, key: str, content_type_key: str, default_content_type: str): async def handle_upload(prompt_id: str, data, key: str, content_type_key: str, default_content_type: str):
items = data.get(key, []) items = data.get(key, [])
# upload_tasks = [] upload_tasks = []
for item in items: for item in items:
# Skipping temp files # Skipping temp files
@ -1375,33 +1375,41 @@ async def handle_upload(prompt_id: str, data, key: str, content_type_key: str, d
elif file_extension == '.webp': elif file_extension == '.webp':
file_type = 'image/webp' file_type = 'image/webp'
# upload_tasks.append(upload_file( upload_tasks.append(upload_file(
# prompt_id,
# item.get("filename"),
# subfolder=item.get("subfolder"),
# type=item.get("type"),
# content_type=file_type,
# item=item
# ))
await upload_file(
prompt_id, prompt_id,
item.get("filename"), item.get("filename"),
subfolder=item.get("subfolder"), subfolder=item.get("subfolder"),
type=item.get("type"), type=item.get("type"),
content_type=file_type, content_type=file_type,
item=item item=item
) ))
# await upload_file(
# prompt_id,
# item.get("filename"),
# subfolder=item.get("subfolder"),
# type=item.get("type"),
# content_type=file_type,
# item=item
# )
# Execute all upload tasks concurrently # Execute all upload tasks concurrently
# await asyncio.gather(*upload_tasks) await asyncio.gather(*upload_tasks)
# Upload files in the background # Upload files in the background
async def upload_in_background(prompt_id: str, data, node_id=None, have_upload=True, node_meta=None): async def upload_in_background(prompt_id: str, data, node_id=None, have_upload=True, node_meta=None):
try: try:
await handle_upload(prompt_id, data, 'images', "content_type", "image/png") # await handle_upload(prompt_id, data, 'images', "content_type", "image/png")
await handle_upload(prompt_id, data, 'files', "content_type", "image/png") # await handle_upload(prompt_id, data, 'files', "content_type", "image/png")
await handle_upload(prompt_id, data, 'gifs', "format", "image/gif") # await handle_upload(prompt_id, data, 'gifs', "format", "image/gif")
await handle_upload(prompt_id, data, 'mesh', "format", "application/octet-stream") # await handle_upload(prompt_id, data, 'mesh', "format", "application/octet-stream")
upload_tasks = [
handle_upload(prompt_id, data, 'images', "content_type", "image/png"),
handle_upload(prompt_id, data, 'files', "content_type", "image/png"),
handle_upload(prompt_id, data, 'gifs', "format", "image/gif"),
handle_upload(prompt_id, data, 'mesh', "format", "application/octet-stream")
]
await asyncio.gather(*upload_tasks)
status_endpoint = prompt_metadata[prompt_id].status_endpoint status_endpoint = prompt_metadata[prompt_id].status_endpoint
token = prompt_metadata[prompt_id].token token = prompt_metadata[prompt_id].token