From 80413a7628a86376b0d990e6d792b92e5a954404 Mon Sep 17 00:00:00 2001 From: "Timothy J. Baek" Date: Tue, 30 Apr 2024 11:52:08 -0700 Subject: [PATCH] refac --- backend/apps/images/main.py | 23 +++++++++++------------ 1 file changed, 11 insertions(+), 12 deletions(-) diff --git a/backend/apps/images/main.py b/backend/apps/images/main.py index 8c5457e8a..88cecc940 100644 --- a/backend/apps/images/main.py +++ b/backend/apps/images/main.py @@ -320,16 +320,19 @@ def save_b64_image(b64_str): header, encoded = b64_str.split(",", 1) mime_type = header.split(";")[0] - image_format = mimetypes.guess_extension(mime_type) img_data = base64.b64decode(encoded) + 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: f.write(img_data) - return image_id, image_format + return image_filename except Exception as e: log.exception(f"Error saving image: {e}") - return None, None + return None def save_url_image(url): @@ -395,10 +398,8 @@ def generate_image( images = [] for image in res["data"]: - image_id, image_format = save_b64_image(image["b64_json"]) - images.append( - {"url": f"/cache/image/generations/{image_id}{image_format}"} - ) + image_filename = save_b64_image(image["b64_json"]) + images.append({"url": f"/cache/image/generations/{image_filename}"}) file_body_path = IMAGE_CACHE_DIR.joinpath(f"{image_id}.json") with open(file_body_path, "w") as f: @@ -474,10 +475,8 @@ def generate_image( images = [] for image in res["images"]: - image_id, image_format = save_b64_image(image) - images.append( - {"url": f"/cache/image/generations/{image_id}{image_format}"} - ) + image_filename = save_b64_image(image) + images.append({"url": f"/cache/image/generations/{image_filename}"}) file_body_path = IMAGE_CACHE_DIR.joinpath(f"{image_id}.json") with open(file_body_path, "w") as f: