feat: provide option to run the handler locally as API

blib-la · Feb 11, 2024 · 69848e5 · 69848e5
1 parent 1789028
commit 69848e5
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 4 deletions.
diff --git a/docker-compose.yml b/docker-compose.yml
@@ -2,12 +2,13 @@ version: "3.8"
 
 services:
   comfyui:
-    image: timpietruskyblibla/runpod-worker-comfy:1.0.0
+    image: timpietruskyblibla/runpod-worker-comfy:latest
     container_name: comfyui-worker
     environment:
       - NVIDIA_VISIBLE_DEVICES=all
-    env_file:
-      - .env
+      - SERVE_API_LOCALLY=true
+    ports:
+      - "8000:8000"
     runtime: nvidia
     volumes:
       - ./data/comfyui/output:/comfyui/output
diff --git a/src/start.sh b/src/start.sh
@@ -8,4 +8,10 @@ echo "runpod-worker-comfy: Starting ComfyUI"
 python3 /comfyui/main.py --disable-auto-launch --disable-metadata &
 
 echo "runpod-worker-comfy: Starting RunPod Handler"
-python3 -u /rp_handler.py
+
+# Serve the API and don't shutdown the container
+if [ "$SERVE_API_LOCALLY" == "true" ]; then
+    python3 -u /rp_handler.py --rp_serve_api
+else
+    python3 -u /rp_handler.py
+fi