refactor: clean up commented-out code in upload completion logic

* Fix preview image loading and add support for 3D model uploads

* skip: sending the upload status, make sure, we are also sending the success status after all files is uploaded

comfy-nodes/external_audio.py

@@ -1,14 +1,13 @@
-import os
 import io
-import torchaudio
 from folder_paths import get_annotated_filepath
 
+
 class ComfyUIDeployExternalAudio:
     RETURN_TYPES = ("AUDIO",)
     RETURN_NAMES = ("audio",)
     FUNCTION = "load_audio"
     CATEGORY = "🔗ComfyDeploy"
-    
+
     @classmethod
     def INPUT_TYPES(cls):
         return {
@@ -29,30 +28,55 @@ class ComfyUIDeployExternalAudio:
                     "STRING",
                     {"multiline": False, "default": ""},
                 ),
-            }
+            },
         }
-    
+
     @classmethod
     def VALIDATE_INPUTS(s, audio_file, **kwargs):
         return True
 
-    def load_audio(self, input_id, audio_file, default_value=None, display_name=None, description=None):
-        if audio_file and audio_file != "":
-            if audio_file.startswith(('http://', 'https://')):
-                # Handle URL input
-                import requests
-                response = requests.get(audio_file)
-                audio_data = io.BytesIO(response.content)
-                waveform, sample_rate = torchaudio.load(audio_data)
+    def load_audio(
+        self,
+        input_id,
+        audio_file,
+        default_value=None,
+        display_name=None,
+        description=None,
+    ):
+        try:
+            import torchaudio
+
+            if audio_file and audio_file != "":
+                if audio_file.startswith(("http://", "https://")):
+                    # Handle URL input
+                    try:
+                        import requests
+
+                        response = requests.get(audio_file)
+                        audio_data = io.BytesIO(response.content)
+                        waveform, sample_rate = torchaudio.load(audio_data)
+                    except Exception as e:
+                        print(f"Error loading audio from URL: {e}")
+                        return (default_value,)
+                else:
+                    # Handle local file
+                    try:
+                        audio_path = get_annotated_filepath(audio_file)
+                        waveform, sample_rate = torchaudio.load(audio_path)
+                    except Exception as e:
+                        print(f"Error loading local audio file: {e}")
+                        return (default_value,)
+
+                audio = {"waveform": waveform.unsqueeze(0), "sample_rate": sample_rate}
+                return (audio,)
             else:
-                # Handle local file
-                audio_path = get_annotated_filepath(audio_file)
-                waveform, sample_rate = torchaudio.load(audio_path)
-            
-            audio = {"waveform": waveform.unsqueeze(0), "sample_rate": sample_rate}
-            return (audio,)
-        else:
+                return (default_value,)
+        except ImportError as e:
+            print(f"Error: torchaudio not installed or cannot be imported: {e}")
             return (default_value,)
 
+
 NODE_CLASS_MAPPINGS = {"ComfyUIDeployExternalAudio": ComfyUIDeployExternalAudio}
-NODE_DISPLAY_NAME_MAPPINGS = {"ComfyUIDeployExternalAudio": "External Audio (ComfyUI Deploy)"}
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "ComfyUIDeployExternalAudio": "External Audio (ComfyUI Deploy)"
+}

comfy-nodes/output_text.py

@@ -0,0 +1,99 @@
+import os
+import json
+import folder_paths
+
+
+class ComfyDeployOutputText:
+    def __init__(self):
+        self.output_dir = folder_paths.get_output_directory()
+        self.type = "output"
+        self.prefix_append = ""
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "text": (
+                    "STRING",
+                    {
+                        "multiline": True,
+                        "forceInput": True,
+                        "tooltip": "The text to save.",
+                    },
+                ),
+                "filename_prefix": (
+                    "STRING",
+                    {
+                        "default": "ComfyUI",
+                        "tooltip": "The prefix for the file to save. This may include formatting information such as %date:yyyy-MM-dd% to include values from nodes.",
+                    },
+                ),
+                "file_type": (["txt", "json", "md"], {"default": "txt"}),
+            },
+            "optional": {
+                "output_id": (
+                    "STRING",
+                    {"multiline": False, "default": "output_text"},
+                ),
+            },
+            "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},
+        }
+
+    RETURN_TYPES = ()
+    FUNCTION = "run"
+
+    OUTPUT_NODE = True
+    CATEGORY = "🔗ComfyDeploy"
+    DESCRIPTION = "Saves the input text to your ComfyUI output directory."
+
+    def run(
+        self,
+        text,
+        filename_prefix="ComfyUI",
+        file_type="txt",
+        output_id="output_text",
+        prompt=None,
+        extra_pnginfo=None,
+    ):
+        filename_prefix += self.prefix_append
+        # For text, we don't need dimensions, so pass 0, 0
+        full_output_folder, filename, counter, subfolder, filename_prefix = (
+            folder_paths.get_save_image_path(filename_prefix, self.output_dir, 0, 0)
+        )
+
+        results = list()
+
+        # Create file path
+        file = f"{filename}_{counter:05}_.{file_type}"
+        file_path = os.path.join(full_output_folder, file)
+
+        # Save the text based on file type
+        if file_type == "json":
+            try:
+                # Try to save as JSON if the text is valid JSON
+                json_data = json.loads(text) if isinstance(text, str) else text
+                with open(file_path, "w", encoding="utf-8") as f:
+                    json.dump(json_data, f, indent=2)
+            except json.JSONDecodeError:
+                # Fall back to saving as plain text if not valid JSON
+                with open(file_path, "w", encoding="utf-8") as f:
+                    f.write(text)
+        else:
+            # Save as plain text for txt and md
+            with open(file_path, "w", encoding="utf-8") as f:
+                f.write(text)
+
+        results.append(
+            {
+                "filename": file,
+                "subfolder": subfolder,
+                "type": self.type,
+                "output_id": output_id,
+            }
+        )
+
+        return {"ui": {"text_file": results}}
+
+
+NODE_CLASS_MAPPINGS = {"ComfyDeployOutputText": ComfyDeployOutputText}
+NODE_DISPLAY_NAME_MAPPINGS = {"ComfyDeployOutputText": "Text Output (ComfyDeploy)"}

custom_routes.py

@@ -139,7 +139,9 @@ async def async_request_with_retry(
                 logger.error(f"Error response body: {error_body}")
 
             if attempt == max_retries - 1:
-                logger.error(f"Request {method} : {url} failed after {max_retries} attempts: {e}")
+                logger.error(
+                    f"Request {method} : {url} failed after {max_retries} attempts: {e}"
+                )
                 raise
 
         await asyncio.sleep(retry_delay)
@@ -413,6 +415,9 @@ def apply_inputs_to_workflow(workflow_api: Any, inputs: Any, sid: str = None):
                 if value["class_type"] == "ComfyUIDeployExternalImageBatch":
                     value["inputs"]["images"] = new_value
 
+                if value["class_type"] == "ComfyUIDeployExternalEnum":
+                    value["inputs"]["default_value"] = new_value
+
                 if value["class_type"] == "ComfyUIDeployExternalLora":
                     value["inputs"]["lora_url"] = new_value
 
@@ -1310,7 +1315,7 @@ send_json = prompt_server.send_json
 
 
 async def send_json_override(self, event, data, sid=None):
-    # logger.info("INTERNAL:", event, data, sid)
+    # logger.info(f"INTERNAL: event={event}, data={data}, sid={sid}")
     prompt_id = data.get("prompt_id")
 
     target_sid = sid
@@ -1395,9 +1400,12 @@ async def send_json_override(self, event, data, sid=None):
     # the last executing event is none, then the workflow is finished
     if event == "executing" and data.get("node") is None:
         mark_prompt_done(prompt_id=prompt_id)
+        # We will now rely on the UploadQueue worker to set the final SUCCESS status
+        # after all uploads are confirmed complete.
+
         if not have_pending_upload(prompt_id):
-            await update_run(prompt_id, Status.SUCCESS)
-            if prompt_id in prompt_metadata:
+            # await update_run(prompt_id, Status.SUCCESS) # <-- REMOVE/COMMENT OUT
+            if prompt_id in prompt_metadata:  # <-- REMOVE/COMMENT OUT THIS BLOCK
                 current_time = time.perf_counter()
                 if prompt_metadata[prompt_id].start_time is not None:
                     elapsed_time = current_time - prompt_metadata[prompt_id].start_time
@@ -1839,8 +1847,13 @@ async def upload_file(
 
 def have_pending_upload(prompt_id):
     # 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])}")
+    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")
@@ -1946,7 +1959,7 @@ async def handle_upload(
 
     for item in items:
         # Skipping temp files
-        if item.get("type") == "temp":
+        if isinstance(item, dict) and item.get("type") == "temp":
             continue
 
         file_type = item.get(content_type_key, default_content_type)
@@ -1997,6 +2010,7 @@ async def upload_in_background(
                 ("files", "content_type", "image/png"),
                 ("gifs", "format", "image/gif"),
                 ("model_file", "format", "application/octet-stream"),
+                ("text_file", "format", "text/plain"),
             ]:
                 items = data.get(file_type, [])
 
@@ -2005,14 +2019,19 @@ async def upload_in_background(
                     if file_type == "model_file":
                         if isinstance(item, str):
                             filename = os.path.basename(item)
+                            # Extract folder name from the path
+                            folder_path = os.path.dirname(item)
+                            subfolder = (
+                                os.path.basename(folder_path) if folder_path else ""
+                            )
                             item = {
                                 "filename": filename,
-                                "subfolder": "",
+                                "subfolder": subfolder,
                                 "type": "output",
                             }
 
                     # Skip temp files
-                    if item.get("type") == "temp":
+                    if isinstance(item, dict) and item.get("type") == "temp":
                         continue
 
                     # Add to the upload queue instead of uploading immediately
@@ -2077,6 +2096,7 @@ async def update_run_with_output(
             or "files" in data
             or "gifs" in data
             or "model_file" in data
+            or "text_file" in data
         )
     if bypass_upload and have_upload_media:
         print(
@@ -2387,6 +2407,11 @@ class UploadQueue:
             logger.warning(f"No upload endpoint for prompt ID: {prompt_id}")
             return
 
+        # Check if file_info is a valid dictionary with a filename
+        if not isinstance(file_info, dict) or "filename" not in file_info:
+            logger.warning(f"Invalid file_info for prompt ID {prompt_id}: {file_info}")
+            return
+
         filename = file_info.get("filename")
         subfolder = file_info.get("subfolder")
         file_type = file_info.get("type", "output")
@@ -2537,9 +2562,11 @@ class UploadQueue:
                     # If this was the last file for this prompt, show the stats summary
                     if (
                         prompt_id in self.pending_uploads
+                        # We now rely on the worker's finally block for the final SUCCESS update.
+                        # Check if the set becomes empty *after* removal in the worker.
                         and len(self.pending_uploads[prompt_id]) == 1
                     ):
-                        await update_run(prompt_id, Status.SUCCESS)
+                        # await update_run(prompt_id, Status.SUCCESS) # <-- REMOVE/COMMENT OUT
 
                         self._log_upload_stats(prompt_id)
                         # Clean up stats
@@ -2661,6 +2688,8 @@ class UploadQueue:
                 node_id = upload_task["node_id"]
                 upload_id = upload_task["upload_id"]
 
+                print(file_info)
+
                 try:
                     # Coordinate the actual start of the upload
                     async with self.upload_lock: