refac: comfyui

This commit is contained in:
Timothy Jaeryang Baek 2025-01-16 11:17:37 -08:00
parent 5526c43853
commit 7cf4c9c89c

View File

@ -415,10 +415,14 @@ def save_b64_image(b64_str):
return None return None
def save_url_image(url): def save_url_image(url, headers=None):
image_id = str(uuid.uuid4()) image_id = str(uuid.uuid4())
try: try:
if headers:
r = requests.get(url, headers=headers)
else:
r = requests.get(url) r = requests.get(url)
r.raise_for_status() r.raise_for_status()
if r.headers["content-type"].split("/")[0] == "image": if r.headers["content-type"].split("/")[0] == "image":
mime_type = r.headers["content-type"] mime_type = r.headers["content-type"]
@ -542,7 +546,13 @@ async def image_generations(
images = [] images = []
for image in res["data"]: 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}"}) images.append({"url": f"/cache/image/generations/{image_filename}"})
file_body_path = IMAGE_CACHE_DIR.joinpath(f"{image_filename}.json") file_body_path = IMAGE_CACHE_DIR.joinpath(f"{image_filename}.json")