forked from open-webui/open-webui
commit
9df4227db4
1 changed files with 23 additions and 11 deletions
|
@ -317,12 +317,13 @@ class GenerateImageForm(BaseModel):
|
|||
|
||||
def save_b64_image(b64_str):
|
||||
try:
|
||||
image_id = str(uuid.uuid4())
|
||||
|
||||
if "," in b64_str:
|
||||
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}"
|
||||
|
@ -330,6 +331,17 @@ def save_b64_image(b64_str):
|
|||
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
|
||||
|
|
Loading…
Reference in a new issue