370 lines
15 KiB
Python
370 lines
15 KiB
Python
import json
|
||
import copy
|
||
from collections import deque
|
||
from typing import Dict, List, Set, Tuple
|
||
|
||
|
||
from .side_effects import SIDE_EFFECT_NODES
|
||
# SIDE_EFFECT_NODES = {}
|
||
|
||
|
||
class WorkflowGraph:
|
||
"""一个封装了工作流图结构和操作的类。"""
|
||
|
||
def __init__(self, workflow: Dict):
|
||
self.workflow = workflow
|
||
self.adj: Dict[str, List[str]] = {node_id: [] for node_id in workflow}
|
||
self.rev_adj: Dict[str, List[str]] = {node_id: [] for node_id in workflow}
|
||
self.in_degree: Dict[str, int] = {node_id: 0 for node_id in workflow}
|
||
self.out_degree: Dict[str, int] = {node_id: 0 for node_id in workflow}
|
||
self._build()
|
||
|
||
def _build(self):
|
||
"""从工作流字典构建图的邻接表和度信息。"""
|
||
for node_id, node_data in self.workflow.items():
|
||
for input_value in node_data.get("inputs", {}).values():
|
||
if isinstance(input_value, list) and len(input_value) == 2:
|
||
source_id = str(input_value[0])
|
||
if source_id in self.workflow:
|
||
self.add_edge(source_id, node_id)
|
||
|
||
def add_edge(self, source_id: str, dest_id: str):
|
||
"""向图中添加一条边并更新所有相关结构。"""
|
||
if dest_id not in self.adj.get(source_id, []):
|
||
self.adj.setdefault(source_id, []).append(dest_id)
|
||
self.rev_adj.setdefault(dest_id, []).append(source_id)
|
||
self.out_degree[source_id] = self.out_degree.get(source_id, 0) + 1
|
||
self.in_degree[dest_id] = self.in_degree.get(dest_id, 0) + 1
|
||
|
||
def remove_edge(self, source_id: str, dest_id: str):
|
||
"""从图中移除一条边并更新所有相关结构。"""
|
||
if source_id in self.adj and dest_id in self.adj[source_id]:
|
||
self.adj[source_id].remove(dest_id)
|
||
self.out_degree[source_id] -= 1
|
||
if dest_id in self.rev_adj and source_id in self.rev_adj[dest_id]:
|
||
self.rev_adj[dest_id].remove(source_id)
|
||
self.in_degree[dest_id] -= 1
|
||
|
||
def topological_sort(self) -> Tuple[List[str], bool]:
|
||
"""使用 Kahn 算法进行拓扑排序。返回排序后的节点列表和是否成功的标志。"""
|
||
sorted_nodes = []
|
||
q = deque([node_id for node_id, degree in self.in_degree.items() if degree == 0])
|
||
|
||
temp_in_degree = self.in_degree.copy()
|
||
|
||
while q:
|
||
u = q.popleft()
|
||
sorted_nodes.append(u)
|
||
for v in self.adj.get(u, []):
|
||
temp_in_degree[v] -= 1
|
||
if temp_in_degree[v] == 0:
|
||
q.append(v)
|
||
|
||
is_valid = len(sorted_nodes) == len(self.workflow)
|
||
return sorted_nodes, is_valid
|
||
|
||
def get_leaf_nodes(self) -> Set[str]:
|
||
"""获取图中所有的叶子节点（出度为0）。"""
|
||
return {node_id for node_id, degree in self.out_degree.items() if degree == 0}
|
||
|
||
|
||
def _find_function_subgraph(end_node_id: str, graph: WorkflowGraph) -> Set[str]:
|
||
"""从 __FunctionEnd__ 节点开始后向遍历，找到完整的函数子图。"""
|
||
subgraph_nodes = set()
|
||
q = deque([end_node_id])
|
||
visited = {end_node_id}
|
||
while q:
|
||
curr_id = q.popleft()
|
||
subgraph_nodes.add(curr_id)
|
||
for pred_id in graph.rev_adj.get(curr_id, []):
|
||
if pred_id not in visited:
|
||
visited.add(pred_id)
|
||
q.append(pred_id)
|
||
return subgraph_nodes
|
||
|
||
|
||
def _find_param_nodes(subgraph_nodes: Set[str], workflow: Dict, graph: WorkflowGraph) -> Tuple[Set[str], Set[str]]:
|
||
param_source_nodes = {n for n in subgraph_nodes if workflow[n].get("class_type") == "__FunctionParam__"}
|
||
|
||
param_nodes = set()
|
||
q = deque(list(param_source_nodes))
|
||
visited = set(param_source_nodes)
|
||
while q:
|
||
curr_id = q.popleft()
|
||
param_nodes.add(curr_id)
|
||
for succ_id in graph.adj.get(curr_id, []):
|
||
if succ_id in subgraph_nodes and succ_id not in visited:
|
||
visited.add(succ_id)
|
||
q.append(succ_id)
|
||
|
||
return param_nodes
|
||
|
||
|
||
def _transform_single_function(end_node_id: str, workflow: Dict, graph: WorkflowGraph):
|
||
"""转换单个 __FunctionEnd__ 节点及其关联的子图。"""
|
||
subgraph_nodes = _find_function_subgraph(end_node_id, graph)
|
||
original_end_node = workflow[end_node_id]
|
||
enable_capture = original_end_node["inputs"].get("capture", True)
|
||
|
||
if enable_capture:
|
||
param_nodes = _find_param_nodes(subgraph_nodes, workflow, graph)
|
||
param_nodes.add(end_node_id)
|
||
capture_nodes = subgraph_nodes - param_nodes
|
||
else:
|
||
param_nodes = subgraph_nodes
|
||
capture_nodes = set()
|
||
|
||
# 恢复原始节点类型以便深拷贝
|
||
workflow[end_node_id]["class_type"] = original_end_node["class_type"]
|
||
|
||
# 步骤 4: 创建 __CreateClosure__ 节点
|
||
closure_node = copy.deepcopy(original_end_node)
|
||
closure_node["class_type"] = "__CreateClosure__"
|
||
closure_node["_meta"] = {"title": "Create Closure"}
|
||
workflow[end_node_id] = closure_node
|
||
|
||
# 步骤 5 & 6: 处理捕获边和函数体
|
||
has_side_effects = False
|
||
param_subgraph_copy = {}
|
||
capture_edge_map = {}
|
||
capture_counter = 0
|
||
|
||
for p_node_id in param_nodes:
|
||
if workflow[p_node_id]["class_type"] == "__FunctionParam__":
|
||
continue
|
||
if workflow[p_node_id]["class_type"] in SIDE_EFFECT_NODES:
|
||
has_side_effects = True
|
||
|
||
node_copy = copy.deepcopy(workflow[p_node_id])
|
||
|
||
for input_name, input_link in list(node_copy.get("inputs", {}).items()):
|
||
if not (isinstance(input_link, list) and len(input_link) == 2):
|
||
continue
|
||
|
||
source_id_str = str(input_link[0])
|
||
|
||
if source_id_str in capture_nodes:
|
||
source_id, source_idx = source_id_str, input_link[1]
|
||
edge_tuple = (source_id, source_idx)
|
||
|
||
if edge_tuple not in capture_edge_map:
|
||
capture_edge_map[edge_tuple] = capture_counter
|
||
closure_node["inputs"][f"capture_{capture_counter}"] = [source_id, source_idx]
|
||
graph.add_edge(source_id, end_node_id)
|
||
capture_counter += 1
|
||
|
||
capture_id = capture_edge_map[edge_tuple]
|
||
node_copy["inputs"][input_name] = ["__capture", capture_id]
|
||
|
||
elif workflow.get(source_id_str, {}).get("class_type") == "__FunctionParam__":
|
||
param_node = workflow[source_id_str]
|
||
param_index = param_node["inputs"].get("index", 0)
|
||
node_copy["inputs"][input_name] = ["__param", param_index]
|
||
|
||
param_subgraph_copy[p_node_id] = node_copy
|
||
|
||
output_edge = param_subgraph_copy[end_node_id]["inputs"].get("return_value")
|
||
closure_node["inputs"]["output"] = json.dumps(output_edge)
|
||
del param_subgraph_copy[end_node_id]
|
||
closure_node["inputs"]["body"] = json.dumps(param_subgraph_copy)
|
||
closure_node["inputs"]["side_effects"] = float("NaN") if has_side_effects else 0.0
|
||
|
||
if "return_value" in closure_node["inputs"]:
|
||
return_link = closure_node["inputs"]["return_value"]
|
||
if isinstance(return_link, list) and len(return_link) == 2:
|
||
source_id = str(return_link[0])
|
||
graph.remove_edge(source_id, end_node_id)
|
||
del closure_node["inputs"]["return_value"]
|
||
|
||
|
||
def _prune_unreachable_nodes(workflow: Dict, leaf_nodes: Set[str]) -> Dict:
|
||
"""从叶子节点回溯，移除所有不可达的节点。"""
|
||
reachable_nodes = set()
|
||
q = deque(list(leaf_nodes))
|
||
visited = set(leaf_nodes)
|
||
while q:
|
||
curr_id = q.popleft()
|
||
reachable_nodes.add(curr_id)
|
||
|
||
node_data = workflow.get(curr_id)
|
||
if node_data:
|
||
for input_value in node_data.get("inputs", {}).values():
|
||
if isinstance(input_value, list) and len(input_value) == 2:
|
||
pred_id = str(input_value[0])
|
||
if pred_id in workflow and pred_id not in visited:
|
||
visited.add(pred_id)
|
||
q.append(pred_id)
|
||
|
||
return {node_id: node_data for node_id, node_data in workflow.items() if node_id in reachable_nodes}
|
||
|
||
|
||
def _preprocess_inspect_nodes(workflow: Dict, execute_outputs: list[str]):
|
||
"""
|
||
预处理 __Inspect__ 节点。
|
||
查找所有仅由 __Inspect__ 节点连接的叶子节点，将其"吸收"到 __Inspect__ 节点的 body 中，
|
||
然后将 __Inspect__ 节点转换为 __InspectImpl__ 节点。
|
||
此函数会直接在传入的 workflow 字典上进行修改。
|
||
"""
|
||
temp_graph = WorkflowGraph(workflow)
|
||
leaf_nodes = temp_graph.get_leaf_nodes()
|
||
|
||
nodes_to_remove = set()
|
||
modified_inspect_nodes: Dict[str, Dict] = {}
|
||
|
||
for leaf_id in leaf_nodes:
|
||
parents = temp_graph.rev_adj.get(leaf_id, [])
|
||
# 条件：叶子节点，且只有一个父节点，且父节点类型为 __Inspect__
|
||
if len(parents) == 1:
|
||
parent_id = parents[0]
|
||
parent_node = workflow.get(parent_id)
|
||
|
||
if not (parent_node and parent_node.get("class_type") == "__Inspect__"):
|
||
continue
|
||
|
||
# 新增检查：确认叶子节点连接的是 __Inspect__ 的 value 输出（索引 1）
|
||
is_connected_to_value_output = False
|
||
leaf_node = workflow[leaf_id]
|
||
for input_link in leaf_node.get("inputs", {}).values():
|
||
if isinstance(input_link, list) and str(input_link[0]) == parent_id and input_link[1] == 1:
|
||
is_connected_to_value_output = True
|
||
break
|
||
|
||
# 如果不是连接到 value 输出，则跳过此叶子节点
|
||
if not is_connected_to_value_output:
|
||
continue
|
||
|
||
# --- 确认无误后，开始转换逻辑 ---
|
||
if parent_id not in modified_inspect_nodes:
|
||
modified_inspect_nodes[parent_id] = {}
|
||
|
||
body = modified_inspect_nodes[parent_id]
|
||
leaf_node_copy = copy.deepcopy(leaf_node)
|
||
|
||
# 查找并替换来自父节点的输入边
|
||
for input_name, input_link in leaf_node_copy.get("inputs", {}).items():
|
||
if isinstance(input_link, list) and str(input_link[0]) == parent_id and input_link[1] == 1:
|
||
leaf_node_copy["inputs"][input_name] = ["__value", 0]
|
||
break
|
||
|
||
body[leaf_id] = leaf_node_copy
|
||
nodes_to_remove.add(leaf_id)
|
||
|
||
# 将更改应用到工作流
|
||
# 1. 更新被修改的 __Inspect__ 节点
|
||
for inspect_id, body_dict in modified_inspect_nodes.items():
|
||
inspect_node = workflow[inspect_id]
|
||
inspect_node["class_type"] = "__InspectImpl__"
|
||
inspect_node["inputs"]["body"] = json.dumps(body_dict)
|
||
|
||
# 2. 移除被吸收的叶子节点
|
||
for node_id in nodes_to_remove:
|
||
if node_id in workflow:
|
||
del workflow[node_id]
|
||
# Remove from execute_outputs if present
|
||
if node_id in execute_outputs:
|
||
execute_outputs.remove(node_id)
|
||
|
||
|
||
def transform_workflow(workflow: dict, execute_outputs: list[str]) -> tuple[dict, list]:
|
||
"""
|
||
根据指定算法变换 ComfyUI 工作流。
|
||
此算法首先预处理 __Inspect__ 节点以支持动态输出，然后将由
|
||
__FunctionParam__ 和 __FunctionEnd__ 定义的函数子图转换为一个 __CreateClosure__ 节点。
|
||
|
||
Args:
|
||
workflow: 代表 ComfyUI 工作流的字典对象。
|
||
|
||
Returns:
|
||
一个元组，包含：
|
||
- 变换后的工作流字典。
|
||
- 一个包含警告信息的列表（如果有）。
|
||
"""
|
||
warnings = []
|
||
workflow_og = workflow
|
||
workflow = copy.deepcopy(workflow_og)
|
||
|
||
# 步骤 2: 预处理 __Inspect__ 节点
|
||
_preprocess_inspect_nodes(workflow, execute_outputs)
|
||
|
||
# 步骤 3: 从可能已修改的工作流重建图结构以进行后续操作
|
||
graph = WorkflowGraph(workflow)
|
||
sorted_nodes, is_valid_dag = graph.topological_sort()
|
||
if not is_valid_dag:
|
||
warnings.append("Warning: The workflow contains cycles and cannot be fully sorted topologically.")
|
||
return workflow_og, warnings
|
||
|
||
# 在函数变换之前，确定最终用于剪枝的叶子节点集合
|
||
leaf_nodes_for_pruning = graph.get_leaf_nodes()
|
||
|
||
# 步骤 4: 按拓扑顺序迭代，转换所有函数定义
|
||
for node_id in sorted_nodes:
|
||
if workflow[node_id].get("class_type") == "__FunctionEnd__":
|
||
_transform_single_function(node_id, workflow, graph)
|
||
|
||
# 步骤 5: 剪枝，删除不再被任何叶子节点依赖的节点
|
||
pruned_workflow = _prune_unreachable_nodes(workflow, leaf_nodes_for_pruning)
|
||
|
||
# 步骤 6: 最终检查，确保图中不包含 __FunctionParam__ 节点
|
||
for node_id, node_data in pruned_workflow.items():
|
||
if node_data.get("class_type") == "__FunctionParam__":
|
||
warnings.append(
|
||
f"Warning: __FunctionParam__ node '{node_id}' remains in the workflow. This indicates incomplete function definitions."
|
||
)
|
||
|
||
return pruned_workflow, warnings
|
||
|
||
|
||
if __name__ == "__main__":
|
||
# 示例 ComfyUI 工作流 JSON 对象 (结构简化)
|
||
# A -> B, A -> C, B -> C, C -> D, D -> E
|
||
# C 是 __FunctionParam__, E 是 __FunctionEnd__
|
||
# A 是捕获变量, B/C/D/E 是函数体
|
||
example_workflow = {
|
||
"1": { # A: Load Checkpoint (capture)
|
||
"inputs": {},
|
||
"class_type": "CheckpointLoaderSimple·",
|
||
"_meta": {"title": "Load Checkpoint"},
|
||
},
|
||
"2": { # B: CLIP Text Encode (param)
|
||
"inputs": {"clip": ["1", 1]},
|
||
"class_type": "CLIPTextEncode",
|
||
"_meta": {"title": "Prompt"},
|
||
},
|
||
"3": { # C: __FunctionParam__ (param source)
|
||
"inputs": {"index": 0},
|
||
"class_type": "__FunctionParam__",
|
||
"_meta": {"title": "Function Param: latent_image"},
|
||
},
|
||
"4": { # D: KSampler (param)
|
||
"inputs": {
|
||
"model": ["1", 0], # Input from capture node "1"
|
||
"positive": ["2", 0],
|
||
"latent_image": ["3", 0],
|
||
},
|
||
"class_type": "KSampler",
|
||
"_meta": {"title": "KSampler"},
|
||
},
|
||
"5": { # E: __FunctionEnd__ -> __CreateClosure__
|
||
"inputs": {"return_value": ["4", 0]},
|
||
"class_type": "__FunctionEnd__",
|
||
"_meta": {"title": "Function End"},
|
||
},
|
||
"6": { # F: Final Output Node
|
||
"inputs": {"images": ["5", 0]},
|
||
"class_type": "SaveImage",
|
||
"_meta": {"title": "Save Image"},
|
||
},
|
||
}
|
||
|
||
# 调用变换函数
|
||
transformed_workflow, warnings_list = transform_workflow(example_workflow)
|
||
|
||
# 打印结果
|
||
print("--- Transformed Workflow ---")
|
||
print(json.dumps(transformed_workflow, indent=2))
|
||
|
||
if warnings_list:
|
||
print("\n--- Warnings ---")
|
||
for warning in warnings_list:
|
||
print(warning)
|