Timothy J. Baek před 1 rokem
rodič
revize
80413a7628
1 změnil soubory, kde provedl 11 přidání a 12 odebrání
  1. 11 12
      backend/apps/images/main.py

+ 11 - 12
backend/apps/images/main.py

@@ -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"])
-                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")
                 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)
-                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")
                 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: