From 9045237bfbd74e2c975604e19fdbddb608d85fda Mon Sep 17 00:00:00 2001
From: psychedelicious <4822129+psychedelicious@users.noreply.github.com>
Date: Fri, 7 Mar 2025 08:38:13 +1000
Subject: [PATCH] feat(api): add util to extract metadata from image

---
 .../app/api/extract_metadata_from_image.py    | 108 ++++++++++++++++++
 1 file changed, 108 insertions(+)
 create mode 100644 invokeai/app/api/extract_metadata_from_image.py

diff --git a/invokeai/app/api/extract_metadata_from_image.py b/invokeai/app/api/extract_metadata_from_image.py
new file mode 100644
index 0000000000..5153b524f2
--- /dev/null
+++ b/invokeai/app/api/extract_metadata_from_image.py
@@ -0,0 +1,108 @@
+import json
+import logging
+from dataclasses import dataclass
+
+from PIL import Image
+
+from invokeai.app.services.shared.graph import Graph
+from invokeai.app.services.workflow_records.workflow_records_common import WorkflowWithoutIDValidator
+
+
+@dataclass
+class ExtractedMetadata:
+    invokeai_metadata: str | None
+    invokeai_workflow: str | None
+    invokeai_graph: str | None
+
+
+def extract_metadata_from_image(
+    pil_image: Image.Image,
+    invokeai_metadata_override: str | None,
+    invokeai_workflow_override: str | None,
+    invokeai_graph_override: str | None,
+    logger: logging.Logger,
+) -> ExtractedMetadata:
+    """
+    Extracts the "invokeai_metadata", "invokeai_workflow", and "invokeai_graph" data embedded in the PIL Image.
+
+    These items are stored as stringified JSON in the image file's metadata, so we need to do some parsing to validate
+    them. Once parsed, the values are returned as they came (as strings), or None if they are not present or invalid.
+
+    In some situations, we may prefer to override the values extracted from the image file with some other values.
+
+    For example, when uploading an image via API, the client can optionally provide the metadata directly in the request,
+    as opposed to embedding it in the image file. In this case, the client-provided metadata will be used instead of the
+    metadata embedded in the image file.
+
+    Args:
+        pil_image: The PIL Image object.
+        invokeai_metadata_override: The metadata override provided by the client.
+        invokeai_workflow_override: The workflow override provided by the client.
+        invokeai_graph_override: The graph override provided by the client.
+        logger: The logger to use for debug logging.
+
+    Returns:
+        ExtractedMetadata: The extracted metadata, workflow, and graph.
+    """
+
+    # The fallback value for metadata is None.
+    stringified_metadata: str | None = None
+
+    # Use the metadata override if provided, else attempt to extract it from the image file.
+    metadata_raw = invokeai_metadata_override or pil_image.info.get("invokeai_metadata", None)
+
+    # If the metadata is present in the image file, we will attempt to parse it as JSON. When we create images,
+    # we always store metadata as a stringified JSON dict. So, we expect it to be a string here.
+    if isinstance(metadata_raw, str):
+        try:
+            # Must be a JSON string
+            metadata_parsed = json.loads(metadata_raw)
+            # Must be a dict
+            if isinstance(metadata_parsed, dict):
+                # Looks good, overwrite the fallback value
+                stringified_metadata = metadata_raw
+        except Exception as e:
+            logger.debug(f"Failed to parse metadata for uploaded image, {e}")
+            pass
+
+    # We expect the workflow, if embedded in the image, to be a JSON-stringified WorkflowWithoutID. We will store it
+    # as a string.
+    workflow_raw: str | None = invokeai_workflow_override or pil_image.info.get("invokeai_workflow", None)
+
+    # The fallback value for workflow is None.
+    stringified_workflow: str | None = None
+
+    # If the workflow is present in the image file, we will attempt to parse it as JSON. When we create images, we
+    # always store workflows as a stringified JSON WorkflowWithoutID. So, we expect it to be a string here.
+    if isinstance(workflow_raw, str):
+        try:
+            # Validate the workflow JSON before storing it
+            WorkflowWithoutIDValidator.validate_json(workflow_raw)
+            # Looks good, overwrite the fallback value
+            stringified_workflow = workflow_raw
+        except Exception:
+            logger.debug("Failed to parse workflow for uploaded image")
+            pass
+
+    # We expect the workflow, if embedded in the image, to be a JSON-stringified Graph. We will store it as a
+    # string.
+    graph_raw: str | None = invokeai_graph_override or pil_image.info.get("invokeai_graph", None)
+
+    # The fallback value for graph is None.
+    stringified_graph: str | None = None
+
+    # If the graph is present in the image file, we will attempt to parse it as JSON. When we create images, we
+    # always store graphs as a stringified JSON Graph. So, we expect it to be a string here.
+    if isinstance(graph_raw, str):
+        try:
+            # Validate the graph JSON before storing it
+            Graph.model_validate_json(graph_raw)
+            # Looks good, overwrite the fallback value
+            stringified_graph = graph_raw
+        except Exception as e:
+            logger.debug(f"Failed to parse graph for uploaded image, {e}")
+            pass
+
+    return ExtractedMetadata(
+        invokeai_metadata=stringified_metadata, invokeai_workflow=stringified_workflow, invokeai_graph=stringified_graph
+    )