fix: image gen

This commit is contained in:
Timothy J. Baek 2024-05-02 15:54:31 -07:00
parent 1f98c091f0
commit 482f010d22

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)