fix(api): deduplicate metadata/workflow extraction logic

This commit is contained in:
psychedelicious 2023-10-18 20:10:17 +11:00
parent 23fa2e560a
commit 2faed653d7
2 changed files with 5 additions and 15 deletions

View File

@ -60,10 +60,10 @@ async def upload_image(
bbox = pil_image.getbbox() bbox = pil_image.getbbox()
pil_image = pil_image.crop(bbox) pil_image = pil_image.crop(bbox)
except Exception: except Exception:
# Error opening the image
ApiDependencies.invoker.services.logger.error(traceback.format_exc()) ApiDependencies.invoker.services.logger.error(traceback.format_exc())
raise HTTPException(status_code=415, detail="Failed to read image") raise HTTPException(status_code=415, detail="Failed to read image")
# TODO: retain non-invokeai metadata on upload?
# attempt to parse metadata from image # attempt to parse metadata from image
metadata_raw = pil_image.info.get("invokeai_metadata", None) metadata_raw = pil_image.info.get("invokeai_metadata", None)
if metadata_raw: if metadata_raw:

View File

@ -65,20 +65,10 @@ class DiskImageFileStorage(ImageFileStorageBase):
pnginfo = PngImagePlugin.PngInfo() pnginfo = PngImagePlugin.PngInfo()
if metadata is not None or workflow is not None: if metadata is not None:
if metadata is not None: pnginfo.add_text("invokeai_metadata", metadata.model_dump_json())
pnginfo.add_text("invokeai_metadata", metadata.model_dump_json()) if workflow is not None:
if workflow is not None: pnginfo.add_text("invokeai_workflow", workflow.model_dump_json())
pnginfo.add_text("invokeai_workflow", workflow.model_dump_json())
else:
# For uploaded images, we want to retain metadata. PIL strips it on save; manually add it back
# TODO: retain non-invokeai metadata on save...
original_metadata = image.info.get("invokeai_metadata", None)
if original_metadata is not None:
pnginfo.add_text("invokeai_metadata", original_metadata)
original_workflow = image.info.get("invokeai_workflow", None)
if original_workflow is not None:
pnginfo.add_text("invokeai_workflow", original_workflow)
image.save( image.save(
image_path, image_path,