From 236a174b18e1599310ebc28ab831c457df07bdc4 Mon Sep 17 00:00:00 2001
From: Bahadir Ciloglu <bahadirciloglu@gmail.com>
Date: Sat, 1 Nov 2025 16:02:20 +0300
Subject: [PATCH] Rewrite runpod_handler.py with clean formatting - no
 whitespace issues

---
 ComfyUI-master/runpod_handler.py | 43 +++++++++++++++++---------------
 1 file changed, 23 insertions(+), 20 deletions(-)

diff --git a/ComfyUI-master/runpod_handler.py b/ComfyUI-master/runpod_handler.py
index fae07d1f3..d0cc90000 100644
--- a/ComfyUI-master/runpod_handler.py
+++ b/ComfyUI-master/runpod_handler.py
@@ -22,6 +22,7 @@ logging.basicConfig(
 )
 logger = logging.getLogger(__name__)
 
+
 class ComfyUIServerlessHandler:
     def __init__(self):
         self.comfyui_url = "http://127.0.0.1:8000"
@@ -41,8 +42,8 @@ class ComfyUIServerlessHandler:
         try:
             logger.info("Starting ComfyUI server...")
             self.comfyui_process = subprocess.Popen([
-                "python3", "main.py", 
-                "--listen", "0.0.0.0", 
+                "python3", "main.py",
+                "--listen", "0.0.0.0",
                 "--port", "8000",
                 "--dont-print-server",
                 "--disable-auto-launch"
@@ -64,31 +65,31 @@ class ComfyUIServerlessHandler:
             except requests.exceptions.RequestException:
                 pass
             time.sleep(2)
-        
+
         logger.error(f"ComfyUI not ready after {timeout} seconds")
         return False
 
     def download_input_files(self, input_data: Dict[str, Any]) -> Dict[str, str]:
         """Download input files and return local paths"""
         local_files = {}
-        
+
         if "input_files" in input_data:
             for file_key, file_url in input_data["input_files"].items():
                 try:
                     response = requests.get(file_url, timeout=60)
                     response.raise_for_status()
-                    
+
                     # Create temporary file
                     with tempfile.NamedTemporaryFile(
-                        delete=False, 
+                        delete=False,
                         dir="/tmp/inputs",
                         suffix=os.path.splitext(file_url)[1]
                     ) as tmp_file:
                         tmp_file.write(response.content)
                         local_files[file_key] = tmp_file.name
-                        
+
                     logger.info(f"Downloaded {file_key} to {local_files[file_key]}")
-                    
+
                 except Exception as e:
                     logger.error(f"Failed to download {file_key}: {str(e)}")
                     raise
@@ -108,10 +109,10 @@ class ComfyUIServerlessHandler:
                 timeout=30
             )
             queue_response.raise_for_status()
-            
+
             prompt_id = queue_response.json()["prompt_id"]
             logger.info(f"Queued workflow with prompt_id: {prompt_id}")
-            
+
             # Wait for completion
             return self.wait_for_completion(prompt_id)
 
@@ -122,17 +123,17 @@ class ComfyUIServerlessHandler:
     def wait_for_completion(self, prompt_id: str, timeout: int = 300) -> Dict[str, Any]:
         """Wait for workflow completion and return results"""
         start_time = time.time()
-        
+
         while time.time() - start_time < timeout:
             try:
                 # Check queue status
                 queue_response = requests.get(f"{self.comfyui_url}/queue")
                 queue_data = queue_response.json()
-                
+
                 # Check if our job is still in queue
                 running = any(item[1]["prompt_id"] == prompt_id for item in queue_data.get("queue_running", []))
                 pending = any(item[1]["prompt_id"] == prompt_id for item in queue_data.get("queue_pending", []))
-                
+
                 if not running and not pending:
                     # Job completed, get results
                     history_response = requests.get(f"{self.comfyui_url}/history/{prompt_id}")
@@ -140,9 +141,9 @@ class ComfyUIServerlessHandler:
                         history_data = history_response.json()
                         if prompt_id in history_data:
                             return self.process_results(history_data[prompt_id])
-                
+
                 time.sleep(2)
-                
+
             except Exception as e:
                 logger.error(f"Error checking completion: {str(e)}")
                 time.sleep(5)
@@ -156,7 +157,7 @@ class ComfyUIServerlessHandler:
             "outputs": [],
             "metadata": {}
         }
-        
+
         if "outputs" in history_data:
             for node_id, node_output in history_data["outputs"].items():
                 if "images" in node_output:
@@ -168,23 +169,23 @@ class ComfyUIServerlessHandler:
                             "subfolder": image_info.get("subfolder", ""),
                             "type": image_info.get("type", "output")
                         }
-                        
+
                         try:
                             image_response = requests.get(image_url, params=params)
                             image_response.raise_for_status()
-                            
+
                             # Save to temp file for upload
                             output_path = f"/tmp/outputs/{image_info['filename']}"
                             with open(output_path, "wb") as f:
                                 f.write(image_response.content)
-                            
+
                             results["outputs"].append({
                                 "type": "image",
                                 "filename": image_info["filename"],
                                 "path": output_path,
                                 "node_id": node_id
                             })
-                            
+
                         except Exception as e:
                             logger.error(f"Failed to process image {image_info['filename']}: {str(e)}")
 
@@ -202,6 +203,7 @@ class ComfyUIServerlessHandler:
         except Exception as e:
             logger.error(f"Cleanup failed: {str(e)}")
 
+
 def handler(job: Dict[str, Any]) -> Dict[str, Any]:
     """Main serverless handler function"""
     handler_instance = ComfyUIServerlessHandler()
@@ -260,6 +262,7 @@ def handler(job: Dict[str, Any]) -> Dict[str, Any]:
         # Always cleanup
         handler_instance.cleanup()
 
+
 if __name__ == "__main__":
     # Start the serverless worker
     runpod.serverless.start({"handler": handler})
\ No newline at end of file