feat: added env vars COMFY_POLLING_INTERVAL_MS and COMFY_POLLING_MAX_…

…RETRIES
blib-la · Mar 19, 2024 · 18c510c · 18c510c
1 parent 4b48ef0
commit 18c510c
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 8 deletions.
diff --git a/README.md b/README.md
@@ -19,9 +19,9 @@ Read our article here: https://blib.la/blog/comfyui-on-runpod
 - [Config](#config)
   - [Upload image to AWS S3](#upload-image-to-aws-s3)
 - [Use the Docker image on RunPod](#use-the-docker-image-on-runpod)
-  * [Bring your own models](#bring-your-own-models)
-    + [Network Volume](#network-volume)
-    + [Custom Docker Image](#custom-docker-image)
+  - [Bring your own models](#bring-your-own-models)
+    - [Network Volume](#network-volume)
+    - [Custom Docker Image](#custom-docker-image)
 - [API specification](#api-specification)
   - [JSON Request Body](#json-request-body)
   - [Fields](#fields)
@@ -69,9 +69,11 @@ Read our article here: https://blib.la/blog/comfyui-on-runpod
 
 ## Config
 
-| Environment Variable | Description                                                                                                                                                                        | Default |
-| -------------------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ------- |
-| `REFRESH_WORKER`     | When you want stop the worker after each finished job to have a clean state, see [official documentation](https://docs.runpod.io/docs/handler-additional-controls#refresh-worker). | `false` |
+| Environment Variable        | Description                                                                                                                                                                           | Default |
+| --------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ------- |
+| `REFRESH_WORKER`            | When you want to stop the worker after each finished job to have a clean state, see [official documentation](https://docs.runpod.io/docs/handler-additional-controls#refresh-worker). | `false` |
+| `COMFY_POLLING_INTERVAL_MS` | Time to wait between poll attempts in milliseconds.                                                                                                                                   | `250`   |
+| `COMFY_POLLING_MAX_RETRIES` | Maximum number of poll attempts. You have to increase when you have long running workflows.                                                                                           | `500`   |
 
 ### Upload image to AWS S3
 

diff --git a/src/rp_handler.py b/src/rp_handler.py
@@ -14,9 +14,9 @@
 # Maximum number of API check attempts
 COMFY_API_AVAILABLE_MAX_RETRIES = 500
 # Time to wait between poll attempts in milliseconds
-COMFY_POLLING_INTERVAL_MS = 250
+COMFY_POLLING_INTERVAL_MS = os.environ.get("COMFY_POLLING_INTERVAL_MS", 250)
 # Maximum number of poll attempts
-COMFY_POLLING_MAX_RETRIES = 500
+COMFY_POLLING_MAX_RETRIES = os.environ.get("COMFY_POLLING_MAX_RETRIES", 500)
 # Host where ComfyUI is running
 COMFY_HOST = "127.0.0.1:8188"
 # Enforce a clean state after each job is done