From 313ce956fd0f7d6ce7c519a7fdd889de7a3ca372 Mon Sep 17 00:00:00 2001 From: BennyKok Date: Mon, 21 Apr 2025 14:08:56 +0800 Subject: [PATCH 1/2] Fix preview image loading and add support for 3D model uploads --- custom_routes.py | 23 +++++++---------------- 1 file changed, 7 insertions(+), 16 deletions(-) diff --git a/custom_routes.py b/custom_routes.py index 37e30f0..19035c8 100644 --- a/custom_routes.py +++ b/custom_routes.py @@ -1838,13 +1838,9 @@ async def upload_file( def have_pending_upload(prompt_id): - if ( - prompt_id in prompt_metadata - and len(prompt_metadata[prompt_id].uploading_nodes) > 0 - ): - logger.info( - f"Have pending upload {len(prompt_metadata[prompt_id].uploading_nodes)}" - ) + # Check if there are pending uploads in the queue + if prompt_id in upload_queue.pending_uploads and upload_queue.pending_uploads[prompt_id]: + logger.info(f"Have pending upload {len(upload_queue.pending_uploads[prompt_id])}") return True logger.info("No pending upload") @@ -1898,16 +1894,11 @@ async def handle_error(prompt_id, data, e: Exception): async def update_file_status( prompt_id: str, data, uploading, have_error=False, node_id=None ): - # if 'uploading_nodes' not in prompt_metadata[prompt_id]: - # prompt_metadata[prompt_id]['uploading_nodes'] = set() + # We're using upload_queue as the single source of truth for tracking uploads + # The upload_queue.pending_uploads is managed by the UploadQueue class itself + # We no longer need to track uploading_nodes in prompt_metadata - if node_id is not None: - if uploading: - prompt_metadata[prompt_id].uploading_nodes.add(node_id) - else: - prompt_metadata[prompt_id].uploading_nodes.discard(node_id) - - # logger.info(f"Remaining uploads: {prompt_metadata[prompt_id].uploading_nodes}") + # logger.info(f"Pending uploads in queue: {upload_queue.pending_uploads.get(prompt_id, set())}") # Update the remote status if have_error: From 9beac36d0ffaa4d2f20d797b020b577dc1fe6a7b Mon Sep 17 00:00:00 2001 From: BennyKok Date: Mon, 21 Apr 2025 14:18:18 +0800 Subject: [PATCH 2/2] skip: sending the upload status, make sure, we are also sending the success status after all files is uploaded --- custom_routes.py | 20 +++++++++++--------- 1 file changed, 11 insertions(+), 9 deletions(-) diff --git a/custom_routes.py b/custom_routes.py index 19035c8..02d4317 100644 --- a/custom_routes.py +++ b/custom_routes.py @@ -1912,15 +1912,15 @@ async def update_file_status( return # if there are still nodes that are uploading, then we set the status to uploading - if uploading: - if prompt_metadata[prompt_id].status != Status.UPLOADING: - await update_run(prompt_id, Status.UPLOADING) - await send( - "uploading", - { - "prompt_id": prompt_id, - }, - ) + # if uploading: + # if prompt_metadata[prompt_id].status != Status.UPLOADING: + # await update_run(prompt_id, Status.UPLOADING) + # await send( + # "uploading", + # { + # "prompt_id": prompt_id, + # }, + # ) # if there are no nodes that are uploading, then we set the status to success elif ( @@ -2539,6 +2539,8 @@ class UploadQueue: prompt_id in self.pending_uploads and len(self.pending_uploads[prompt_id]) == 1 ): + await update_run(prompt_id, Status.SUCCESS) + self._log_upload_stats(prompt_id) # Clean up stats del self.upload_stats[prompt_id]