60 lines
1.9 KiB
Python
60 lines
1.9 KiB
Python
from .classes import ComfyUIWorkflowGraph, ComfyUIWorkflowGraphNode, ComfyUIWorkflowGraphLink, LLMGeneratedGraph, _get_widget_values
|
|
from .comfy_types import input_output_types, get_type_from_input_name
|
|
from nodes import NODE_CLASS_MAPPINGS
|
|
|
|
|
|
def transform_input_to_graph(llm_graph) -> ComfyUIWorkflowGraph:
|
|
llm_graph = LLMGeneratedGraph(
|
|
nodes=llm_graph["nodes"],
|
|
edges=llm_graph["edges"]
|
|
)
|
|
last_node_id = llm_graph.get_last_node_id()
|
|
last_link_id = llm_graph.get_last_link_id()
|
|
|
|
comfy_graph = ComfyUIWorkflowGraph(
|
|
last_node_id=llm_graph.get_last_node_id(),
|
|
last_link_id=llm_graph.get_last_link_id(),
|
|
)
|
|
|
|
# Add nodes
|
|
for index, node in enumerate(llm_graph.nodes):
|
|
widget_values = _get_widget_values(node)
|
|
new_node = ComfyUIWorkflowGraphNode(
|
|
id=node["id"],
|
|
type=node["class_type"],
|
|
order=index,
|
|
widgets_values=widget_values
|
|
)
|
|
|
|
comfy_graph.add_node(new_node)
|
|
|
|
for index, edge in enumerate(llm_graph.edges):
|
|
link_id = index + 1
|
|
to_node = comfy_graph.get_node_by_id(edge['to'])
|
|
from_node = comfy_graph.get_node_by_id(edge['from'])
|
|
|
|
to_node.add_input(
|
|
edge['to_input'],
|
|
get_type_from_input_name(edge['to_input']),
|
|
link_id,
|
|
)
|
|
from_node.add_output(
|
|
link_id,
|
|
edge['from_output']
|
|
)
|
|
|
|
# Create links
|
|
link = ComfyUIWorkflowGraphLink(
|
|
id=link_id,
|
|
from_node_id=edge['from'],
|
|
from_node_output_slot=edge['from_output'],
|
|
to_node_id=edge['to'],
|
|
to_node_input_slot=to_node._get_input_slot_index_from_name(edge['to_input']),
|
|
from_node_output_type=from_node._get_output_type_from_slot_index(edge['from_output'])
|
|
)
|
|
|
|
comfy_graph.add_link(link)
|
|
|
|
|
|
return comfy_graph
|