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

fix: image gen

Browse source
This commit is contained in:
Timothy J. Baek 2024-05-02 15:54:31 -07:00
parent 1f98c091f0
commit 482f010d22
1 changed files with 9 additions and 9 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -348,18 +348,20 @@ def save_url_image(url):
if not image_format: if not image_format:
raise ValueError("Could not determine image type from MIME type") raise ValueError("Could not determine image type from MIME type")
file_path = IMAGE_CACHE_DIR.joinpath(f"{image_id}{image_format}") image_filename = f"{image_id}{image_format}"
file_path = IMAGE_CACHE_DIR.joinpath(f"{image_filename}")
with open(file_path, "wb") as image_file: with open(file_path, "wb") as image_file:
for chunk in r.iter_content(chunk_size=8192): for chunk in r.iter_content(chunk_size=8192):
image_file.write(chunk) image_file.write(chunk)
return image_id, image_format return image_filename
else: else:
log.error(f"Url does not point to an image.") log.error(f"Url does not point to an image.")
return None, None return None
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
@app.post("/generations") @app.post("/generations")
@ -435,11 +437,9 @@ def generate_image(
images = [] images = []
for image in res["data"]: for image in res["data"]:
image_id, image_format = save_url_image(image["url"]) image_filename = save_url_image(image["url"])
images.append( images.append({"url": f"/cache/image/generations/{image_filename}"})
{"url": f"/cache/image/generations/{image_id}{image_format}"} file_body_path = IMAGE_CACHE_DIR.joinpath(f"{image_filename}.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:
json.dump(data.model_dump(exclude_none=True), f) json.dump(data.model_dump(exclude_none=True), f)