From 683e6ca337985c73222f46753d7442a35363205a Mon Sep 17 00:00:00 2001 From: "Timothy J. Baek" Date: Thu, 2 May 2024 15:59:45 -0700 Subject: [PATCH] fix --- backend/apps/images/main.py | 34 +++++++++++++++++++++++----------- 1 file changed, 23 insertions(+), 11 deletions(-) diff --git a/backend/apps/images/main.py b/backend/apps/images/main.py index eeaa2b40..f45cf0d1 100644 --- a/backend/apps/images/main.py +++ b/backend/apps/images/main.py @@ -317,19 +317,31 @@ class GenerateImageForm(BaseModel): def save_b64_image(b64_str): try: - header, encoded = b64_str.split(",", 1) - mime_type = header.split(";")[0] - - img_data = base64.b64decode(encoded) - image_id = str(uuid.uuid4()) - 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_filename + if "," in b64_str: + header, encoded = b64_str.split(",", 1) + mime_type = header.split(";")[0] + + img_data = base64.b64decode(encoded) + 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_filename + else: + image_filename = f"{image_id}.png" + file_path = IMAGE_CACHE_DIR.joinpath(image_filename) + + img_data = base64.b64decode(b64_str) + + # Write the image data to a file + with open(file_path, "wb") as f: + f.write(img_data) + return image_filename + except Exception as e: log.exception(f"Error saving image: {e}") return None