refactor: 优化工作流队列管理和提示构建逻辑,增加节点排序和过滤功能,提升代码可读性和处理效率

This commit is contained in:
iHeyTang 2025-08-18 19:48:48 +08:00
parent 7484e300ec
commit d66b640d4f
2 changed files with 29 additions and 7 deletions

View File

@ -119,7 +119,9 @@ class WorkflowQueueManager:
time_diff = current_time - self._last_processing_time
minutes = int(time_diff.total_seconds() // 60)
seconds = int(time_diff.total_seconds() % 60)
time_since_last_processing = f"上次处理: {minutes}{seconds}秒前"
time_since_last_processing = (
f"上次处理: {minutes}{seconds}秒前"
)
else:
time_since_last_processing = "上次处理: 从未"
@ -429,7 +431,6 @@ async def _queue_prompt(
"extra_pnginfo": {"workflow": workflow},
},
}
logger.info(f"提交到 ComfyUI /prompt 端点的payload: {json.dumps(payload)}")
async with aiohttp.ClientSession(timeout=ClientTimeout(total=90)) as session:
prompt_url = f"{http_url}/prompt"
try:

View File

@ -1,3 +1,5 @@
from ast import List
import json
import logging
from typing import Any, Dict
@ -22,6 +24,7 @@ async def build_prompt(
"""
构建prompt
"""
patched_workflow = await _patch_workflow(
workflow_data, api_spec, request_data, server
)
@ -269,6 +272,10 @@ async def _patch_workflow(
continue
workflow_data["nodes"] = list(nodes_map.values())
# 对workflow_data中的nodes进行排序按照id从小到大
workflow_data["nodes"] = sorted(workflow_data["nodes"], key=lambda x: x["id"])
return workflow_data
@ -300,7 +307,6 @@ def _convert_workflow_to_prompt_api_format(workflow_data: dict) -> dict:
for node in workflow_data["nodes"]:
node_id = str(node["id"])
inputs_dict = {}
# 1. 处理控件输入 (widgets)
widgets_values = node.get("widgets_values", [])
widget_cursor = 0
@ -313,6 +319,10 @@ def _convert_workflow_to_prompt_api_format(workflow_data: dict) -> dict:
inputs_dict[widget_name] = widgets_values[widget_cursor]
widget_cursor += 1
# 特殊处理:如果节点是 EG_CYQ_JB 则需要忽略widgets_values的第1个值这个是用来控制control net的它不参与赋值
if node["type"] == "EG_CYQ_JB" and widget_cursor == 1:
widget_cursor += 1
# 2. 处理连接输入 (links)
for input_config in node.get("inputs", []):
if "link" in input_config and input_config["link"] is not None:
@ -321,7 +331,18 @@ def _convert_workflow_to_prompt_api_format(workflow_data: dict) -> dict:
# 输入名称是input_config中的'name'
inputs_dict[input_config["name"]] = link_map[link_id]
prompt_api_format[node_id] = {"class_type": node["type"], "inputs": inputs_dict}
prompt_api_format[node_id] = {
"inputs": inputs_dict,
"class_type": node["type"],
"_meta": {
"title": node.get("title", ""),
},
}
# 过滤掉类型为 Note 的节点
prompt_api_format = {
k: v for k, v in prompt_api_format.items() if v["class_type"] not in ["Note"]
}
return prompt_api_format