mirror of
https://github.com/invoke-ai/InvokeAI
synced 2024-08-30 20:32:17 +00:00
56 lines
1.5 KiB
Python
56 lines
1.5 KiB
Python
import json
|
|
from typing import Optional
|
|
|
|
from pydantic import ValidationError
|
|
|
|
from invokeai.app.services.shared.graph import Edge
|
|
|
|
|
|
def get_metadata_graph_from_raw_session(session_raw: str) -> Optional[dict]:
|
|
"""
|
|
Parses raw session string, returning a dict of the graph.
|
|
|
|
Only the general graph shape is validated; none of the fields are validated.
|
|
|
|
Any `metadata_accumulator` nodes and edges are removed.
|
|
|
|
Any validation failure will return None.
|
|
"""
|
|
|
|
graph = json.loads(session_raw).get("graph", None)
|
|
|
|
# sanity check make sure the graph is at least reasonably shaped
|
|
if (
|
|
not isinstance(graph, dict)
|
|
or "nodes" not in graph
|
|
or not isinstance(graph["nodes"], dict)
|
|
or "edges" not in graph
|
|
or not isinstance(graph["edges"], list)
|
|
):
|
|
# something has gone terribly awry, return an empty dict
|
|
return None
|
|
|
|
try:
|
|
# delete the `metadata_accumulator` node
|
|
del graph["nodes"]["metadata_accumulator"]
|
|
except KeyError:
|
|
# no accumulator node, all good
|
|
pass
|
|
|
|
# delete any edges to or from it
|
|
for i, edge in enumerate(graph["edges"]):
|
|
try:
|
|
# try to parse the edge
|
|
Edge(**edge)
|
|
except ValidationError:
|
|
# something has gone terribly awry, return an empty dict
|
|
return None
|
|
|
|
if (
|
|
edge["source"]["node_id"] == "metadata_accumulator"
|
|
or edge["destination"]["node_id"] == "metadata_accumulator"
|
|
):
|
|
del graph["edges"][i]
|
|
|
|
return graph
|