From a26f4306a4613825f1f0ec845a8b89a80cabcba3 Mon Sep 17 00:00:00 2001
From: Peter De-Ath <peter@de-ath.co.uk>
Date: Thu, 29 Aug 2024 21:57:39 +0100
Subject: [PATCH] fix: add useragent default headers comfyui

---
 backend/apps/images/utils/comfyui.py | 25 +++++++++++++++++++++----
 1 file changed, 21 insertions(+), 4 deletions(-)

diff --git a/backend/apps/images/utils/comfyui.py b/backend/apps/images/utils/comfyui.py
index 158484223..fa4736131 100644
--- a/backend/apps/images/utils/comfyui.py
+++ b/backend/apps/images/utils/comfyui.py
@@ -15,20 +15,33 @@ from pydantic import BaseModel
 
 from typing import Optional
 
+default_headers = {"User-Agent": "Mozilla/5.0"}
+
 
 def queue_prompt(prompt, client_id, base_url):
     log.info("queue_prompt")
     p = {"prompt": prompt, "client_id": client_id}
     data = json.dumps(p).encode("utf-8")
-    req = urllib.request.Request(f"{base_url}/prompt", data=data)
-    return json.loads(urllib.request.urlopen(req).read())
+    log.debug(f"queue_prompt data: {data}")
+    try:
+        req = urllib.request.Request(
+            f"{base_url}/prompt", data=data, headers=default_headers
+        )
+        response = urllib.request.urlopen(req).read()
+        return json.loads(response)
+    except Exception as e:
+        log.exception(f"Error while queuing prompt: {e}")
+        raise e
 
 
 def get_image(filename, subfolder, folder_type, base_url):
     log.info("get_image")
     data = {"filename": filename, "subfolder": subfolder, "type": folder_type}
     url_values = urllib.parse.urlencode(data)
-    with urllib.request.urlopen(f"{base_url}/view?{url_values}") as response:
+    req = urllib.request.Request(
+        f"{base_url}/view?{url_values}", headers=default_headers
+    )
+    with urllib.request.urlopen(req) as response:
         return response.read()
 
 
@@ -41,7 +54,11 @@ def get_image_url(filename, subfolder, folder_type, base_url):
 
 def get_history(prompt_id, base_url):
     log.info("get_history")
-    with urllib.request.urlopen(f"{base_url}/history/{prompt_id}") as response:
+
+    req = urllib.request.Request(
+        f"{base_url}/history/{prompt_id}", headers=default_headers
+    )
+    with urllib.request.urlopen(req) as response:
         return json.loads(response.read())