177 lines
6.1 KiB
Python
177 lines
6.1 KiB
Python
from __future__ import annotations
|
|
|
|
from typing import Any
|
|
|
|
from .models import NodeRecord
|
|
|
|
|
|
class GraphIndex:
|
|
def __init__(
|
|
self,
|
|
prompt: dict[str, Any] | None,
|
|
workflow: dict[str, Any] | None,
|
|
) -> None:
|
|
self._workflow_nodes = self._normalize_workflow_nodes(workflow)
|
|
self._nodes = self._normalize_prompt_nodes(prompt)
|
|
if not self._nodes:
|
|
self._nodes = dict(self._workflow_nodes)
|
|
else:
|
|
self._attach_workflow_supplements()
|
|
|
|
def get_node(self, node_id: str | int) -> NodeRecord | None:
|
|
return self._nodes.get(str(node_id))
|
|
|
|
def nodes(self) -> list[NodeRecord]:
|
|
return list(self._nodes.values())
|
|
|
|
def is_link(self, value: Any) -> bool:
|
|
return self.link_target(value) is not None
|
|
|
|
def link_target(self, value: Any) -> tuple[str, int] | None:
|
|
if not isinstance(value, (list, tuple)) or len(value) != 2:
|
|
return None
|
|
node_id, output_index = value
|
|
if not isinstance(node_id, (str, int)):
|
|
return None
|
|
if not isinstance(output_index, int):
|
|
return None
|
|
return str(node_id), output_index
|
|
|
|
def upstream_link_inputs(self, node_id: str | int) -> list[tuple[str, Any]]:
|
|
node = self.get_node(node_id)
|
|
if node is None:
|
|
return []
|
|
return [
|
|
(field_name, value)
|
|
for field_name, value in node.inputs.items()
|
|
if self.is_link(value)
|
|
]
|
|
|
|
def downstream_nodes(self, node_id: str | int) -> list[NodeRecord]:
|
|
source_id = str(node_id)
|
|
downstream: list[NodeRecord] = []
|
|
for node in self._nodes.values():
|
|
for value in node.inputs.values():
|
|
target = self.link_target(value)
|
|
if target and target[0] == source_id:
|
|
downstream.append(node)
|
|
break
|
|
return downstream
|
|
|
|
def workflow_cache_for(
|
|
self,
|
|
node_id: str | int,
|
|
field_names: list[str],
|
|
) -> str | None:
|
|
cache_entry = self.workflow_cache_entry(node_id, field_names)
|
|
if cache_entry is None:
|
|
return None
|
|
value, _source = cache_entry
|
|
return value
|
|
|
|
def workflow_cache_entry(
|
|
self,
|
|
node_id: str | int,
|
|
field_names: list[str],
|
|
) -> tuple[str, str] | None:
|
|
node = self._workflow_nodes.get(str(node_id))
|
|
if node is None:
|
|
return None
|
|
|
|
for field_name in field_names:
|
|
value = node.inputs.get(field_name)
|
|
if isinstance(value, str) and value:
|
|
return value, f"workflow input {field_name}"
|
|
|
|
widgets = node.widgets_values
|
|
if isinstance(widgets, dict):
|
|
for field_name in field_names:
|
|
value = widgets.get(field_name)
|
|
if isinstance(value, str) and value:
|
|
return value, f"workflow widget {field_name}"
|
|
elif isinstance(widgets, list):
|
|
for value in widgets:
|
|
if isinstance(value, str) and value:
|
|
return value, "workflow widget cache fallback"
|
|
return None
|
|
|
|
def workflow_output_name(self, node_id: str | int, output_index: int) -> str | None:
|
|
node = self._workflow_nodes.get(str(node_id))
|
|
if node is None or not isinstance(node.outputs, list):
|
|
return None
|
|
if output_index < 0 or output_index >= len(node.outputs):
|
|
return None
|
|
output = node.outputs[output_index]
|
|
if isinstance(output, str):
|
|
return output
|
|
if not isinstance(output, dict):
|
|
return None
|
|
for field in ("name", "label", "display_name", "displayName"):
|
|
value = output.get(field)
|
|
if isinstance(value, str) and value:
|
|
return value
|
|
return None
|
|
|
|
def _normalize_prompt_nodes(
|
|
self,
|
|
prompt: dict[str, Any] | None,
|
|
) -> dict[str, NodeRecord]:
|
|
if not isinstance(prompt, dict):
|
|
return {}
|
|
|
|
nodes: dict[str, NodeRecord] = {}
|
|
for raw_id, raw_node in prompt.items():
|
|
if not isinstance(raw_node, dict):
|
|
continue
|
|
node_id = str(raw_id)
|
|
class_type = str(raw_node.get("class_type") or raw_node.get("type") or "")
|
|
inputs = raw_node.get("inputs")
|
|
nodes[node_id] = NodeRecord(
|
|
node_id=node_id,
|
|
class_type=class_type,
|
|
inputs=dict(inputs) if isinstance(inputs, dict) else {},
|
|
widgets_values=None,
|
|
outputs=raw_node.get("outputs"),
|
|
)
|
|
return nodes
|
|
|
|
def _normalize_workflow_nodes(
|
|
self,
|
|
workflow: dict[str, Any] | None,
|
|
) -> dict[str, NodeRecord]:
|
|
if not isinstance(workflow, dict):
|
|
return {}
|
|
|
|
raw_nodes = workflow.get("nodes")
|
|
if not isinstance(raw_nodes, list):
|
|
return {}
|
|
|
|
nodes: dict[str, NodeRecord] = {}
|
|
for raw_node in raw_nodes:
|
|
if not isinstance(raw_node, dict):
|
|
continue
|
|
raw_id = raw_node.get("id")
|
|
if raw_id is None:
|
|
continue
|
|
node_id = str(raw_id)
|
|
class_type = str(raw_node.get("class_type") or raw_node.get("type") or "")
|
|
inputs = raw_node.get("inputs")
|
|
nodes[node_id] = NodeRecord(
|
|
node_id=node_id,
|
|
class_type=class_type,
|
|
inputs=dict(inputs) if isinstance(inputs, dict) else {},
|
|
widgets_values=raw_node.get("widgets_values"),
|
|
outputs=raw_node.get("outputs"),
|
|
)
|
|
return nodes
|
|
|
|
def _attach_workflow_supplements(self) -> None:
|
|
for node_id, workflow_node in self._workflow_nodes.items():
|
|
prompt_node = self._nodes.get(node_id)
|
|
if prompt_node is None:
|
|
continue
|
|
if prompt_node.widgets_values is None:
|
|
prompt_node.widgets_values = workflow_node.widgets_values
|
|
if prompt_node.outputs is None:
|
|
prompt_node.outputs = workflow_node.outputs
|