106 lines
4.0 KiB
Python
106 lines
4.0 KiB
Python
# Copyright (c) 2024 Allan Saddi <allan@saddi.com>
|
|
class PromptUtils:
|
|
def __init__(self, prompt):
|
|
self.prompt = prompt
|
|
|
|
def get_node_info(self, node_id: int | str) -> dict:
|
|
node_id = str(node_id)
|
|
return self.prompt[node_id]
|
|
|
|
def is_input_connected(self, node_id: int | str, name: str) -> bool:
|
|
node_info = self.get_node_info(node_id)
|
|
input = node_info.get("inputs", {}).get(name)
|
|
# In reality, input will be a 2-elem list of source node ID + slot
|
|
# if it is connected.
|
|
return input is not None
|
|
|
|
@property
|
|
def output_map(self):
|
|
if hasattr(self, "_output_map"):
|
|
return self._output_map
|
|
# A map of (src node ID, slot) -> (dest node ID, "input name")
|
|
self._output_map: dict[tuple[str, int], list[tuple[str, str]]] = {}
|
|
for dest_id, dest_info in self.prompt.items():
|
|
for input_name, input in dest_info.get("inputs", {}).items():
|
|
if isinstance(input, list) and len(input) == 2:
|
|
# This is apparently a link
|
|
src_id, src_slot = input
|
|
out_key = (src_id, src_slot)
|
|
outputs = self._output_map.get(out_key)
|
|
if outputs is None:
|
|
self._output_map[out_key] = outputs = []
|
|
outputs.append((dest_id, input_name))
|
|
return self._output_map
|
|
|
|
def is_output_connected(self, node_id: int | str, slot: int) -> bool:
|
|
out_key = (str(node_id), slot)
|
|
return out_key in self.output_map
|
|
|
|
# The assumption is that the given output slot can only be connected
|
|
# to nodes of the same type as the one being checked.
|
|
def get_downstream_nodes(self, node_id: int | str, slot: int) -> list[str]:
|
|
downstream: list[str] = []
|
|
|
|
to_check = [str(node_id)]
|
|
while to_check:
|
|
check_id = to_check.pop(0)
|
|
|
|
# Yes, consider the starting node downstream from itself as well
|
|
downstream.append(check_id)
|
|
|
|
outputs = self.output_map.get((check_id, slot))
|
|
if outputs is not None:
|
|
# NB There can be multiple outputs
|
|
out_ids: list[str] = [link_info[0] for link_info in outputs]
|
|
to_check.extend(out_ids)
|
|
|
|
return downstream
|
|
|
|
|
|
class WorkflowUtils:
|
|
def __init__(self, extra_pnginfo):
|
|
# Note: API mode does not have extra_pnginfo (it is None)
|
|
self.extra_pnginfo = extra_pnginfo
|
|
|
|
def get_node_info(self, node_id: str | int) -> dict | None:
|
|
# Be somewhat paranoid since API mode offers much less access
|
|
if (
|
|
self.extra_pnginfo is not None
|
|
and (workflow := self.extra_pnginfo.get("workflow")) is not None
|
|
and (nodes := workflow.get("nodes")) is not None
|
|
):
|
|
node_id = str(node_id)
|
|
# TODO For the time being, we aren't caching anything.
|
|
# Concievably, if there are a large number of nodes, a lot of time
|
|
# could be wasted searching.
|
|
my_node = [node_info for node_info in nodes if node_info["id"] == node_id]
|
|
if my_node:
|
|
return my_node[0]
|
|
return None
|
|
|
|
def set_property(self, node_id: str | int, property_name: str, value):
|
|
if (my_node := self.get_node_info(node_id)) is not None:
|
|
my_node["properties"][property_name] = value
|
|
|
|
def set_widget(self, node_id: str | int, slot: int, value):
|
|
if (my_node := self.get_node_info(node_id)) is not None:
|
|
values = my_node.get("widgets_values", [])
|
|
# TODO When would there be a mismatch?
|
|
if slot < len(values):
|
|
values[slot] = value
|
|
|
|
|
|
class AnyType(str):
|
|
def __ne__(self, other):
|
|
if self == "*":
|
|
# If we're the wildcard, match anything
|
|
return False
|
|
return super().__ne__(other)
|
|
|
|
|
|
class ComboType(str):
|
|
def __ne__(self, other):
|
|
if self == "*" and isinstance(other, list):
|
|
return False
|
|
return True
|