tdpeuter/open-webui
tdpeuter/open-webui
1
Fork 0
forked from open-webui/open-webui
Code Issues Pull requests 1 Releases 2 Packages Activity Actions

refac

Browse source
This commit is contained in:
Timothy J. Baek 2024-04-30 11:52:08 -07:00
parent 3321a1b922
commit 80413a7628
1 changed files with 11 additions and 12 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

23
backend/apps/images/main.py
View file

@ -320,16 +320,19 @@ def save_b64_image(b64_str):
header, encoded = b64_str.split(",", 1) header, encoded = b64_str.split(",", 1)
mime_type = header.split(";")[0] mime_type = header.split(";")[0]
image_format = mimetypes.guess_extension(mime_type)
img_data = base64.b64decode(encoded) img_data = base64.b64decode(encoded)
image_id = str(uuid.uuid4()) image_id = str(uuid.uuid4())
file_path = IMAGE_CACHE_DIR / f"{image_id}{image_format}" image_format = mimetypes.guess_extension(mime_type)
image_filename = f"{image_id}{image_format}"
file_path = IMAGE_CACHE_DIR / f"{image_filename}"
with open(file_path, "wb") as f: with open(file_path, "wb") as f:
f.write(img_data) f.write(img_data)
return image_id, image_format return image_filename
except Exception as e: except Exception as e:
log.exception(f"Error saving image: {e}") log.exception(f"Error saving image: {e}")
return None, None return None
def save_url_image(url): def save_url_image(url):
@ -395,10 +398,8 @@ def generate_image(
images = [] images = []
for image in res["data"]: for image in res["data"]:
image_id, image_format = save_b64_image(image["b64_json"]) image_filename = save_b64_image(image["b64_json"])
images.append( images.append({"url": f"/cache/image/generations/{image_filename}"})
{"url": f"/cache/image/generations/{image_id}{image_format}"}
)
file_body_path = IMAGE_CACHE_DIR.joinpath(f"{image_id}.json") file_body_path = IMAGE_CACHE_DIR.joinpath(f"{image_id}.json")
with open(file_body_path, "w") as f: with open(file_body_path, "w") as f:
@ -474,10 +475,8 @@ def generate_image(
images = [] images = []
for image in res["images"]: for image in res["images"]:
image_id, image_format = save_b64_image(image) image_filename = save_b64_image(image)
images.append( images.append({"url": f"/cache/image/generations/{image_filename}"})
{"url": f"/cache/image/generations/{image_id}{image_format}"}
)
file_body_path = IMAGE_CACHE_DIR.joinpath(f"{image_id}.json") file_body_path = IMAGE_CACHE_DIR.joinpath(f"{image_id}.json")
with open(file_body_path, "w") as f: with open(file_body_path, "w") as f: