From dfb856b92556fd64a99ce9fe28bf608998cf42ae Mon Sep 17 00:00:00 2001 From: Calvin Huang Date: Mon, 2 Dec 2024 12:11:51 +0800 Subject: [PATCH] feat: add retry to queue workflows (cherry picked from commit f0f14dfdecfcbfbd560647e6c70e11f94dc9e71f) --- src/rp_handler.py | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/src/rp_handler.py b/src/rp_handler.py index 04c3b714..b009d0ed 100644 --- a/src/rp_handler.py +++ b/src/rp_handler.py @@ -168,8 +168,17 @@ def queue_workflow(workflow): # The top level element "prompt" is required by ComfyUI data = json.dumps({"prompt": workflow}).encode("utf-8") - req = urllib.request.Request(f"http://{COMFY_HOST}/prompt", data=data) - return json.loads(urllib.request.urlopen(req).read()) + retries = 0 + while retries < COMFY_API_AVAILABLE_MAX_RETRIES: + try: + req = urllib.request.Request(f"http://{COMFY_HOST}/prompt", data=data) + return json.loads(urllib.request.urlopen(req).read()) + except Exception as e: + print(f"Error queuing workflow: {str(e)}. Retrying...") + time.sleep(COMFY_API_AVAILABLE_INTERVAL_MS / 1000) + retries += 1 + + raise Exception("Max retries reached while queuing workflow") def get_history(prompt_id):