From 7cf4c9c89c37c4c3a70b84e1fcdca04337efc080 Mon Sep 17 00:00:00 2001 From: Timothy Jaeryang Baek Date: Thu, 16 Jan 2025 11:17:37 -0800 Subject: [PATCH] refac: comfyui --- backend/open_webui/routers/images.py | 16 +++++++++++++--- 1 file changed, 13 insertions(+), 3 deletions(-) diff --git a/backend/open_webui/routers/images.py b/backend/open_webui/routers/images.py index f4833d0b5..7afd9d106 100644 --- a/backend/open_webui/routers/images.py +++ b/backend/open_webui/routers/images.py @@ -415,10 +415,14 @@ def save_b64_image(b64_str): return None -def save_url_image(url): +def save_url_image(url, headers=None): image_id = str(uuid.uuid4()) try: - r = requests.get(url) + if headers: + r = requests.get(url, headers=headers) + else: + r = requests.get(url) + r.raise_for_status() if r.headers["content-type"].split("/")[0] == "image": mime_type = r.headers["content-type"] @@ -542,7 +546,13 @@ async def image_generations( images = [] for image in res["data"]: - image_filename = save_url_image(image["url"]) + headers = None + if request.app.state.config.COMFYUI_API_KEY: + headers = { + "Authorization": f"Bearer {request.app.state.config.COMFYUI_API_KEY}" + } + + image_filename = save_url_image(image["url"], headers) images.append({"url": f"/cache/image/generations/{image_filename}"}) file_body_path = IMAGE_CACHE_DIR.joinpath(f"{image_filename}.json")