fix: add useragent default headers comfyui

This commit is contained in:
Peter De-Ath 2024-08-29 21:57:39 +01:00
parent 6589450bd6
commit a26f4306a4

View File

@ -15,20 +15,33 @@ from pydantic import BaseModel
from typing import Optional from typing import Optional
default_headers = {"User-Agent": "Mozilla/5.0"}
def queue_prompt(prompt, client_id, base_url): def queue_prompt(prompt, client_id, base_url):
log.info("queue_prompt") log.info("queue_prompt")
p = {"prompt": prompt, "client_id": client_id} p = {"prompt": prompt, "client_id": client_id}
data = json.dumps(p).encode("utf-8") data = json.dumps(p).encode("utf-8")
req = urllib.request.Request(f"{base_url}/prompt", data=data) log.debug(f"queue_prompt data: {data}")
return json.loads(urllib.request.urlopen(req).read()) 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): def get_image(filename, subfolder, folder_type, base_url):
log.info("get_image") log.info("get_image")
data = {"filename": filename, "subfolder": subfolder, "type": folder_type} data = {"filename": filename, "subfolder": subfolder, "type": folder_type}
url_values = urllib.parse.urlencode(data) 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() return response.read()
@ -41,7 +54,11 @@ def get_image_url(filename, subfolder, folder_type, base_url):
def get_history(prompt_id, base_url): def get_history(prompt_id, base_url):
log.info("get_history") 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()) return json.loads(response.read())