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

Merge pull request #1930 from open-webui/dev

Browse source 
fix
This commit is contained in:
Timothy Jaeryang Baek 2024-05-02 16:00:05 -07:00 committed by GitHub
commit 9df4227db4
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 23 additions and 11 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -317,12 +317,13 @@ class GenerateImageForm(BaseModel):
def save_b64_image(b64_str): def save_b64_image(b64_str):
try: try:
image_id = str(uuid.uuid4())
if "," in b64_str:
header, encoded = b64_str.split(",", 1) header, encoded = b64_str.split(",", 1)
mime_type = header.split(";")[0] mime_type = header.split(";")[0]
img_data = base64.b64decode(encoded) img_data = base64.b64decode(encoded)
image_id = str(uuid.uuid4())
image_format = mimetypes.guess_extension(mime_type) image_format = mimetypes.guess_extension(mime_type)
image_filename = f"{image_id}{image_format}" image_filename = f"{image_id}{image_format}"
@ -330,6 +331,17 @@ def save_b64_image(b64_str):
with open(file_path, "wb") as f: with open(file_path, "wb") as f:
f.write(img_data) f.write(img_data)
return image_filename 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: except Exception as e:
log.exception(f"Error saving image: {e}") log.exception(f"Error saving image: {e}")
return None return None