fix(app): retain graph in metadata when uploading images

This commit is contained in:
psychedelicious 2024-05-17 19:14:59 +10:00
parent 386d552493
commit 93ebc175c6
5 changed files with 16 additions and 6 deletions

View File

@ -49,6 +49,7 @@ async def upload_image(
metadata = None
workflow = None
graph = None
contents = await file.read()
try:
@ -76,9 +77,17 @@ async def upload_image(
try:
workflow = WorkflowWithoutIDValidator.validate_json(workflow_raw)
except ValidationError:
ApiDependencies.invoker.services.logger.warn("Failed to parse metadata for uploaded image")
ApiDependencies.invoker.services.logger.warn("Failed to parse workflow for uploaded image")
pass
# attempt to extract graph from image
graph_raw = pil_image.info.get("invokeai_graph", None)
if isinstance(graph_raw, str):
graph = graph_raw
else:
ApiDependencies.invoker.services.logger.warn("Failed to parse graph for uploaded image")
pass
try:
image_dto = ApiDependencies.invoker.services.images.create(
image=pil_image,
@ -88,6 +97,7 @@ async def upload_image(
board_id=board_id,
metadata=metadata,
workflow=workflow,
graph=graph,
is_intermediate=is_intermediate,
)

View File

@ -36,7 +36,7 @@ class ImageFileStorageBase(ABC):
image_name: str,
metadata: Optional[MetadataField] = None,
workflow: Optional[WorkflowWithoutID] = None,
graph: Optional[Graph] = None,
graph: Optional[Graph | str] = None,
thumbnail_size: int = 256,
) -> None:
"""Saves an image and a 256x256 WEBP thumbnail. Returns a tuple of the image name, thumbnail name, and created timestamp."""

View File

@ -59,7 +59,7 @@ class DiskImageFileStorage(ImageFileStorageBase):
image_name: str,
metadata: Optional[MetadataField] = None,
workflow: Optional[WorkflowWithoutID] = None,
graph: Optional[Graph] = None,
graph: Optional[Graph | str] = None,
thumbnail_size: int = 256,
) -> None:
try:
@ -78,7 +78,7 @@ class DiskImageFileStorage(ImageFileStorageBase):
info_dict["invokeai_workflow"] = workflow_json
pnginfo.add_text("invokeai_workflow", workflow_json)
if graph is not None:
graph_json = graph.model_dump_json()
graph_json = graph.model_dump_json() if isinstance(graph, Graph) else graph
info_dict["invokeai_graph"] = graph_json
pnginfo.add_text("invokeai_graph", graph_json)

View File

@ -54,7 +54,7 @@ class ImageServiceABC(ABC):
is_intermediate: Optional[bool] = False,
metadata: Optional[MetadataField] = None,
workflow: Optional[WorkflowWithoutID] = None,
graph: Optional[Graph] = None,
graph: Optional[Graph | str] = None,
) -> ImageDTO:
"""Creates an image, storing the file and its metadata."""
pass

View File

@ -45,7 +45,7 @@ class ImageService(ImageServiceABC):
is_intermediate: Optional[bool] = False,
metadata: Optional[MetadataField] = None,
workflow: Optional[WorkflowWithoutID] = None,
graph: Optional[Graph] = None,
graph: Optional[Graph | str] = None,
) -> ImageDTO:
if image_origin not in ResourceOrigin:
raise InvalidOriginException