wangbo/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2026-02-10 21:42:37 +08:00
Code Issues Actions 13 Packages Projects Releases Wiki Activity

Rewrite runpod_handler.py with clean formatting - no whitespace issues

Browse Source
This commit is contained in:
Bahadir Ciloglu 2025-11-01 16:02:20 +03:00
parent dc5a696d3b
commit 236a174b18
1 changed files with 23 additions and 20 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

43
ComfyUI-master/runpod_handler.py
View File

@ -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})