Skip to content

Commit 7f04521

Browse files
feat(api): use extract_metadata_from_image in upload router
1 parent 8e46b03 commit 7f04521

File tree

1 file changed

+16
-34
lines changed

1 file changed

+16
-34
lines changed

invokeai/app/api/routers/images.py

Lines changed: 16 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -6,9 +6,10 @@
66
from fastapi.responses import FileResponse
77
from fastapi.routing import APIRouter
88
from PIL import Image
9-
from pydantic import BaseModel, Field, JsonValue
9+
from pydantic import BaseModel, Field
1010

1111
from invokeai.app.api.dependencies import ApiDependencies
12+
from invokeai.app.api.extract_metadata_from_image import extract_metadata_from_image
1213
from invokeai.app.invocations.fields import MetadataField
1314
from invokeai.app.services.image_records.image_records_common import (
1415
ImageCategory,
@@ -45,18 +46,16 @@ async def upload_image(
4546
board_id: Optional[str] = Query(default=None, description="The board to add this image to, if any"),
4647
session_id: Optional[str] = Query(default=None, description="The session ID associated with this upload, if any"),
4748
crop_visible: Optional[bool] = Query(default=False, description="Whether to crop the image"),
48-
metadata: Optional[JsonValue] = Body(
49-
default=None, description="The metadata to associate with the image", embed=True
49+
metadata: Optional[str] = Body(
50+
default=None,
51+
description="The metadata to associate with the image, must be a stringified JSON dict",
52+
embed=True,
5053
),
5154
) -> ImageDTO:
5255
"""Uploads an image"""
5356
if not file.content_type or not file.content_type.startswith("image"):
5457
raise HTTPException(status_code=415, detail="Not an image")
5558

56-
_metadata = None
57-
_workflow = None
58-
_graph = None
59-
6059
contents = await file.read()
6160
try:
6261
pil_image = Image.open(io.BytesIO(contents))
@@ -67,30 +66,13 @@ async def upload_image(
6766
ApiDependencies.invoker.services.logger.error(traceback.format_exc())
6867
raise HTTPException(status_code=415, detail="Failed to read image")
6968

70-
# TODO: retain non-invokeai metadata on upload?
71-
# attempt to parse metadata from image
72-
metadata_raw = metadata if isinstance(metadata, str) else pil_image.info.get("invokeai_metadata", None)
73-
if isinstance(metadata_raw, str):
74-
_metadata = metadata_raw
75-
else:
76-
ApiDependencies.invoker.services.logger.debug("Failed to parse metadata for uploaded image")
77-
pass
78-
79-
# attempt to parse workflow from image
80-
workflow_raw = pil_image.info.get("invokeai_workflow", None)
81-
if isinstance(workflow_raw, str):
82-
_workflow = workflow_raw
83-
else:
84-
ApiDependencies.invoker.services.logger.debug("Failed to parse workflow for uploaded image")
85-
pass
86-
87-
# attempt to extract graph from image
88-
graph_raw = pil_image.info.get("invokeai_graph", None)
89-
if isinstance(graph_raw, str):
90-
_graph = graph_raw
91-
else:
92-
ApiDependencies.invoker.services.logger.debug("Failed to parse graph for uploaded image")
93-
pass
69+
extracted_metadata = extract_metadata_from_image(
70+
pil_image=pil_image,
71+
invokeai_metadata_override=metadata,
72+
invokeai_workflow_override=None,
73+
invokeai_graph_override=None,
74+
logger=ApiDependencies.invoker.services.logger,
75+
)
9476

9577
try:
9678
image_dto = ApiDependencies.invoker.services.images.create(
@@ -99,9 +81,9 @@ async def upload_image(
9981
image_category=image_category,
10082
session_id=session_id,
10183
board_id=board_id,
102-
metadata=_metadata,
103-
workflow=_workflow,
104-
graph=_graph,
84+
metadata=extracted_metadata.invokeai_metadata,
85+
workflow=extracted_metadata.invokeai_workflow,
86+
graph=extracted_metadata.invokeai_graph,
10587
is_intermediate=is_intermediate,
10688
)
10789

0 commit comments

Comments
 (0)