refactor: 移除旧的comfyui_client模块，重构API路由以使用新的comfy_workflow和queue_manager，提升代码结构和可维护性
This commit is contained in:
iHeyTang
443
workflow_service/comfy/comfy_queue.py
Normal file
443
workflow_service/comfy/comfy_queue.py
Normal file
@@ -0,0 +1,443 @@
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import random
|
||||
import uuid
|
||||
import websockets
|
||||
from datetime import datetime
|
||||
from typing import Dict, Any
|
||||
|
||||
import aiohttp
|
||||
from aiohttp import ClientTimeout
|
||||
|
||||
from workflow_service.comfy.comfy_workflow import build_prompt
|
||||
from workflow_service.config import Settings, ComfyUIServer
|
||||
from workflow_service.database.api import (
|
||||
create_workflow_run,
|
||||
update_workflow_run_status,
|
||||
create_workflow_run_nodes,
|
||||
update_workflow_run_node_status,
|
||||
get_workflow_run,
|
||||
get_workflow_run_nodes,
|
||||
)
|
||||
|
||||
settings = Settings()
|
||||
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
API_INPUT_PREFIX = "INPUT_"
|
||||
API_OUTPUT_PREFIX = "OUTPUT_"
|
||||
|
||||
|
||||
class ComfyUIExecutionError(Exception):
|
||||
def __init__(self, error_data: dict):
|
||||
self.error_data = error_data
|
||||
# 创建一个对开发者友好的异常消息
|
||||
message = (
|
||||
f"ComfyUI节点执行失败。节点ID: {error_data.get('node_id')}, "
|
||||
f"节点类型: {error_data.get('node_type')}. "
|
||||
f"错误: {error_data.get('exception_message', 'N/A')}"
|
||||
)
|
||||
super().__init__(message)
|
||||
|
||||
|
||||
# 全局任务队列管理器
|
||||
class WorkflowQueueManager:
|
||||
def __init__(self):
|
||||
self.running_tasks = {} # server_url -> task_info
|
||||
self.pending_tasks = [] # 等待队列
|
||||
self.lock = asyncio.Lock()
|
||||
|
||||
async def add_task(
|
||||
self,
|
||||
workflow_name: str,
|
||||
workflow_data: dict,
|
||||
api_spec: dict,
|
||||
request_data: dict,
|
||||
):
|
||||
"""添加新任务到队列"""
|
||||
workflow_run_id = str(uuid.uuid4())
|
||||
async with self.lock:
|
||||
# 创建任务记录
|
||||
await create_workflow_run(
|
||||
workflow_run_id=workflow_run_id,
|
||||
workflow_name=workflow_name,
|
||||
workflow_json=json.dumps(workflow_data),
|
||||
api_spec=json.dumps(api_spec),
|
||||
request_data=json.dumps(request_data),
|
||||
)
|
||||
|
||||
# 创建工作流节点记录
|
||||
nodes_data = []
|
||||
for node in workflow_data.get("nodes", []):
|
||||
nodes_data.append(
|
||||
{"id": str(node["id"]), "type": node.get("type", "unknown")}
|
||||
)
|
||||
await create_workflow_run_nodes(workflow_run_id, nodes_data)
|
||||
|
||||
# 添加到待处理队列
|
||||
self.pending_tasks.append(workflow_run_id)
|
||||
logger.info(
|
||||
f"任务 {workflow_run_id} 已添加到队列，当前队列长度: {len(self.pending_tasks)}"
|
||||
)
|
||||
|
||||
# 尝试处理队列
|
||||
asyncio.create_task(self._process_queue())
|
||||
|
||||
return workflow_run_id
|
||||
|
||||
async def _process_queue(self):
|
||||
"""处理队列中的任务"""
|
||||
async with self.lock:
|
||||
if not self.pending_tasks:
|
||||
return
|
||||
|
||||
# 检查是否有空闲的服务器
|
||||
available_servers = await self._get_available_servers()
|
||||
if not available_servers:
|
||||
logger.info("没有可用的服务器，等待中...")
|
||||
return
|
||||
|
||||
# 获取一个待处理任务
|
||||
workflow_run_id = self.pending_tasks.pop(0)
|
||||
server = available_servers[0]
|
||||
|
||||
# 标记任务为运行中
|
||||
await update_workflow_run_status(
|
||||
workflow_run_id, "running", server.http_url
|
||||
)
|
||||
self.running_tasks[server.http_url] = {
|
||||
"workflow_run_id": workflow_run_id,
|
||||
"started_at": datetime.now(),
|
||||
}
|
||||
|
||||
# 启动任务执行
|
||||
asyncio.create_task(self._execute_task(workflow_run_id, server))
|
||||
|
||||
async def _get_available_servers(self) -> list[ComfyUIServer]:
|
||||
"""获取可用的服务器"""
|
||||
available_servers = []
|
||||
for server in settings.SERVERS:
|
||||
if server.http_url not in self.running_tasks:
|
||||
# 检查服务器状态
|
||||
try:
|
||||
async with aiohttp.ClientSession() as session:
|
||||
status = await self.get_server_status(server, session)
|
||||
if status["is_reachable"] and status["is_free"]:
|
||||
available_servers.append(server)
|
||||
except Exception as e:
|
||||
logger.warning(f"检查服务器 {server.http_url} 状态时出错: {e}")
|
||||
|
||||
return available_servers
|
||||
|
||||
async def _execute_task(self, workflow_run_id: str, server: ComfyUIServer):
|
||||
"""执行任务"""
|
||||
cleanup_paths = []
|
||||
try:
|
||||
# 获取工作流数据
|
||||
workflow_run = await get_workflow_run(workflow_run_id)
|
||||
if not workflow_run:
|
||||
raise Exception(f"找不到工作流运行记录: {workflow_run_id}")
|
||||
|
||||
workflow_data = json.loads(workflow_run["workflow_json"])
|
||||
api_spec = json.loads(workflow_run["api_spec"])
|
||||
request_data = json.loads(workflow_run["request_data"])
|
||||
|
||||
# 执行工作流
|
||||
result = await _execute_prompt_on_server(
|
||||
workflow_data, api_spec, request_data, server, workflow_run_id
|
||||
)
|
||||
|
||||
# 保存处理后的结果到数据库
|
||||
await update_workflow_run_status(
|
||||
workflow_run_id,
|
||||
"completed",
|
||||
result=json.dumps(result, ensure_ascii=False),
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"执行任务 {workflow_run_id} 时出错: {e}")
|
||||
await update_workflow_run_status(
|
||||
workflow_run_id, "failed", error_message=str(e)
|
||||
)
|
||||
finally:
|
||||
# 清理临时文件
|
||||
if cleanup_paths:
|
||||
logger.info(f"正在清理 {len(cleanup_paths)} 个临时文件...")
|
||||
for path in cleanup_paths:
|
||||
try:
|
||||
if os.path.exists(path):
|
||||
os.remove(path)
|
||||
logger.info(f" - 已删除: {path}")
|
||||
except Exception as e:
|
||||
logger.warning(f" - 删除 {path} 时出错: {e}")
|
||||
|
||||
# 清理运行状态
|
||||
async with self.lock:
|
||||
if server.http_url in self.running_tasks:
|
||||
del self.running_tasks[server.http_url]
|
||||
|
||||
# 继续处理队列
|
||||
asyncio.create_task(self._process_queue())
|
||||
|
||||
async def get_task_status(self, workflow_run_id: str) -> dict:
|
||||
"""获取任务状态"""
|
||||
workflow_run = await get_workflow_run(workflow_run_id)
|
||||
if not workflow_run:
|
||||
return {"error": "任务不存在"}
|
||||
nodes = await get_workflow_run_nodes(workflow_run_id)
|
||||
|
||||
result = {
|
||||
"id": workflow_run_id,
|
||||
"status": workflow_run["status"],
|
||||
"created_at": workflow_run["created_at"],
|
||||
"started_at": workflow_run["started_at"],
|
||||
"completed_at": workflow_run["completed_at"],
|
||||
"server_url": workflow_run["server_url"],
|
||||
"error_message": workflow_run["error_message"],
|
||||
"nodes": nodes,
|
||||
}
|
||||
|
||||
# 如果任务完成，从数据库获取结果
|
||||
if workflow_run["status"] == "completed" and workflow_run.get("result"):
|
||||
try:
|
||||
result["result"] = json.loads(workflow_run["result"])
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
result["result"] = None
|
||||
|
||||
return result
|
||||
|
||||
async def get_server_status(
|
||||
self, server: ComfyUIServer, session: aiohttp.ClientSession
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
检查单个ComfyUI服务器的详细状态。
|
||||
返回一个包含可达性、队列状态和详细队列内容的字典。
|
||||
"""
|
||||
# [BUG修复] 确保初始字典结构与成功时的结构一致，以满足Pydantic模型
|
||||
status_info = {
|
||||
"is_reachable": False,
|
||||
"is_free": False,
|
||||
"queue_details": {"running_count": 0, "pending_count": 0},
|
||||
}
|
||||
try:
|
||||
queue_url = f"{server.http_url}/queue"
|
||||
async with session.get(queue_url, timeout=60) as response:
|
||||
response.raise_for_status()
|
||||
queue_data = await response.json()
|
||||
|
||||
status_info["is_reachable"] = True
|
||||
|
||||
running_count = len(queue_data.get("queue_running", []))
|
||||
pending_count = len(queue_data.get("queue_pending", []))
|
||||
|
||||
status_info["queue_details"] = {
|
||||
"running_count": running_count,
|
||||
"pending_count": pending_count,
|
||||
}
|
||||
|
||||
status_info["is_free"] = running_count == 0 and pending_count == 0
|
||||
|
||||
except Exception as e:
|
||||
# 当请求失败时，将返回上面定义的、结构正确的初始 status_info
|
||||
logger.warning(f"无法检查服务器 {server.http_url} 的队列状态: {e}")
|
||||
|
||||
return status_info
|
||||
|
||||
|
||||
# 全局队列管理器实例
|
||||
queue_manager = WorkflowQueueManager()
|
||||
|
||||
|
||||
async def _execute_prompt_on_server(
|
||||
workflow_data: dict,
|
||||
api_spec: dict,
|
||||
request_data: dict,
|
||||
server: ComfyUIServer,
|
||||
workflow_run_id: str,
|
||||
) -> dict:
|
||||
"""
|
||||
在指定的服务器上执行一个准备好的prompt。
|
||||
现在支持节点级别的状态跟踪。
|
||||
"""
|
||||
client_id = str(uuid.uuid4())
|
||||
|
||||
# 构建prompt
|
||||
prompt = await build_prompt(workflow_data, api_spec, request_data, server)
|
||||
|
||||
# 更新工作流运行状态，记录prompt_id和client_id
|
||||
await update_workflow_run_status(
|
||||
workflow_run_id,
|
||||
"running",
|
||||
server.http_url,
|
||||
None, # prompt_id将在_queue_prompt后更新
|
||||
client_id,
|
||||
)
|
||||
|
||||
# 提交到ComfyUI
|
||||
prompt_id = await _queue_prompt(workflow_data, prompt, client_id, server.http_url)
|
||||
|
||||
# 更新prompt_id
|
||||
await update_workflow_run_status(
|
||||
workflow_run_id, "running", server.http_url, prompt_id, client_id
|
||||
)
|
||||
|
||||
logger.info(
|
||||
f"工作流 {workflow_run_id} 已在 {server.http_url} 上入队，Prompt ID: {prompt_id}"
|
||||
)
|
||||
|
||||
# 获取执行结果，现在支持节点级别的状态跟踪
|
||||
results = await _get_execution_results(
|
||||
workflow_data, prompt_id, client_id, server.ws_url, workflow_run_id
|
||||
)
|
||||
return results
|
||||
|
||||
|
||||
async def _queue_prompt(
|
||||
workflow: dict, prompt: dict, client_id: str, http_url: str
|
||||
) -> str:
|
||||
"""通过HTTP POST将工作流任务提交到指定的ComfyUI服务器。"""
|
||||
for node_id in prompt:
|
||||
prompt[node_id]["inputs"][f"cache_buster_{uuid.uuid4().hex}"] = random.random()
|
||||
payload = {
|
||||
"prompt": prompt,
|
||||
"client_id": client_id,
|
||||
"extra_data": {
|
||||
"api_key_comfy_org": "",
|
||||
"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:
|
||||
async with session.post(prompt_url, json=payload) as response:
|
||||
logger.info(f"ComfyUI /prompt 端点返回的响应: {response}")
|
||||
response.raise_for_status()
|
||||
result = await response.json()
|
||||
logger.info(f"ComfyUI /prompt 端点返回的响应: {result}")
|
||||
if "prompt_id" not in result:
|
||||
raise Exception(f"从 ComfyUI /prompt 端点返回的响应无效: {result}")
|
||||
return result["prompt_id"]
|
||||
except Exception as e:
|
||||
logger.error(f"提交到 ComfyUI /prompt 端点时发生错误: {e}")
|
||||
raise e
|
||||
|
||||
|
||||
async def _get_execution_results(
|
||||
workflow_data: Dict,
|
||||
prompt_id: str,
|
||||
client_id: str,
|
||||
ws_url: str,
|
||||
workflow_run_id: str,
|
||||
) -> dict:
|
||||
"""
|
||||
通过WebSocket连接到指定的ComfyUI服务器，聚合执行结果。
|
||||
现在支持节点级别的状态跟踪。
|
||||
"""
|
||||
full_ws_url = f"{ws_url}?clientId={client_id}"
|
||||
aggregated_outputs = {}
|
||||
|
||||
try:
|
||||
async with websockets.connect(full_ws_url) as websocket:
|
||||
while True:
|
||||
try:
|
||||
out = await websocket.recv()
|
||||
if not isinstance(out, str):
|
||||
continue
|
||||
|
||||
message = json.loads(out)
|
||||
msg_type = message.get("type")
|
||||
data = message.get("data")
|
||||
|
||||
if not (data and data.get("prompt_id") == prompt_id):
|
||||
continue
|
||||
|
||||
# 捕获并处理执行错误
|
||||
if msg_type == "execution_error":
|
||||
error_data = data
|
||||
logger.error(
|
||||
f"ComfyUI执行错误 (Prompt ID: {prompt_id}): {error_data}"
|
||||
)
|
||||
|
||||
# 更新节点状态为失败
|
||||
node_id = error_data.get("node_id")
|
||||
if node_id:
|
||||
await update_workflow_run_node_status(
|
||||
workflow_run_id,
|
||||
node_id,
|
||||
"failed",
|
||||
error_message=error_data.get(
|
||||
"exception_message", "Unknown error"
|
||||
),
|
||||
)
|
||||
|
||||
# 抛出自定义异常，将错误详情传递出去
|
||||
raise ComfyUIExecutionError(error_data)
|
||||
|
||||
# 处理节点开始执行
|
||||
if msg_type == "executing" and data.get("node"):
|
||||
node_id = data.get("node")
|
||||
logger.info(f"节点 {node_id} 开始执行 (Prompt ID: {prompt_id})")
|
||||
|
||||
# 更新节点状态为运行中
|
||||
await update_workflow_run_node_status(
|
||||
workflow_run_id, node_id, "running"
|
||||
)
|
||||
|
||||
# 处理节点执行完成
|
||||
if msg_type == "executed":
|
||||
node_id = data.get("node")
|
||||
output_data = data.get("output")
|
||||
if node_id and output_data:
|
||||
node = next(
|
||||
(
|
||||
x
|
||||
for x in workflow_data["nodes"]
|
||||
if str(x["id"]) == node_id
|
||||
),
|
||||
None,
|
||||
)
|
||||
if (
|
||||
node
|
||||
and node.get("title", "")
|
||||
and node["title"].startswith(API_OUTPUT_PREFIX)
|
||||
):
|
||||
title = node["title"].replace(API_OUTPUT_PREFIX, "")
|
||||
aggregated_outputs[title] = output_data
|
||||
logger.info(
|
||||
f"收到节点 {node_id} 的输出 (Prompt ID: {prompt_id})"
|
||||
)
|
||||
|
||||
# 更新节点状态为完成
|
||||
await update_workflow_run_node_status(
|
||||
workflow_run_id,
|
||||
node_id,
|
||||
"completed",
|
||||
output_data=json.dumps(output_data),
|
||||
)
|
||||
|
||||
# 处理整个工作流执行完成
|
||||
elif msg_type == "executing" and data.get("node") is None:
|
||||
logger.info(f"Prompt ID: {prompt_id} 执行完成。")
|
||||
return aggregated_outputs
|
||||
|
||||
except websockets.exceptions.ConnectionClosed as e:
|
||||
logger.warning(
|
||||
f"WebSocket 连接已关闭 (Prompt ID: {prompt_id})。错误: {e}"
|
||||
)
|
||||
return aggregated_outputs
|
||||
except Exception as e:
|
||||
# 重新抛出我们自己的异常，或者处理其他意外错误
|
||||
if not isinstance(e, ComfyUIExecutionError):
|
||||
logger.error(f"处理 prompt {prompt_id} 时发生意外错误: {e}")
|
||||
raise e
|
||||
|
||||
except websockets.exceptions.InvalidURI as e:
|
||||
logger.error(
|
||||
f"错误: 尝试连接的WebSocket URI无效: '{full_ws_url}'. 原始URL: '{ws_url}'. 错误: {e}"
|
||||
)
|
||||
raise e
|
||||
422
workflow_service/comfy/comfy_workflow.py
Normal file
422
workflow_service/comfy/comfy_workflow.py
Normal file
@@ -0,0 +1,422 @@
|
||||
import logging
|
||||
from typing import Any, Dict
|
||||
|
||||
import aiohttp
|
||||
from workflow_service.config import ComfyUIServer
|
||||
|
||||
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
API_INPUT_PREFIX = "INPUT_"
|
||||
API_OUTPUT_PREFIX = "OUTPUT_"
|
||||
|
||||
|
||||
async def build_prompt(
|
||||
workflow_data: dict,
|
||||
api_spec: dict[str, dict[str, Any]],
|
||||
request_data: dict[str, Any],
|
||||
server: ComfyUIServer,
|
||||
):
|
||||
"""
|
||||
构建prompt
|
||||
"""
|
||||
patched_workflow = await _patch_workflow(
|
||||
workflow_data, api_spec, request_data, server
|
||||
)
|
||||
prompt = _convert_workflow_to_prompt_api_format(patched_workflow)
|
||||
return prompt
|
||||
|
||||
|
||||
def parse_api_spec(workflow_data: dict) -> Dict[str, Dict[str, Any]]:
|
||||
"""
|
||||
解析工作流，并根据规范生成嵌套结构的API参数名。
|
||||
结构: {'节点名': {'node_id': '节点ID', '字段名': {字段信息}}}
|
||||
"""
|
||||
spec = {"inputs": {}, "outputs": {}}
|
||||
if "nodes" not in workflow_data or not isinstance(workflow_data["nodes"], list):
|
||||
raise ValueError(
|
||||
"Invalid workflow format: 'nodes' key not found or is not a list."
|
||||
)
|
||||
|
||||
nodes_map = {str(node["id"]): node for node in workflow_data["nodes"]}
|
||||
|
||||
for node_id, node in nodes_map.items():
|
||||
title: str = node.get("title")
|
||||
if not title:
|
||||
continue
|
||||
|
||||
if title.startswith(API_INPUT_PREFIX):
|
||||
# 提取节点名（去掉API_INPUT_PREFIX前缀）
|
||||
node_name = title[len(API_INPUT_PREFIX) :].lower()
|
||||
|
||||
# 如果节点名不在inputs中，创建新的节点条目
|
||||
if node_name not in spec["inputs"]:
|
||||
spec["inputs"][node_name] = {
|
||||
"node_id": node_id,
|
||||
"class_type": node.get("type"),
|
||||
}
|
||||
|
||||
if "inputs" in node:
|
||||
for a_input in node.get("inputs", []):
|
||||
if a_input.get("link") is None and "widget" in a_input:
|
||||
widget_name = a_input["widget"]["name"].lower()
|
||||
|
||||
# 特殊处理：隐藏LoadImage节点的upload字段
|
||||
if node.get("type") == "LoadImage" and widget_name == "upload":
|
||||
continue
|
||||
|
||||
# 确保inputs字段存在
|
||||
if "inputs" not in spec["inputs"][node_name]:
|
||||
spec["inputs"][node_name]["inputs"] = {}
|
||||
|
||||
# 直接使用widget_name作为字段名
|
||||
spec["inputs"][node_name]["inputs"][widget_name] = {
|
||||
"type": _get_param_type(a_input.get("type", "STRING")),
|
||||
"widget_name": a_input["widget"]["name"],
|
||||
}
|
||||
|
||||
elif title.startswith(API_OUTPUT_PREFIX):
|
||||
# 提取节点名（去掉API_OUTPUT_PREFIX前缀）
|
||||
node_name = title[len(API_OUTPUT_PREFIX) :].lower()
|
||||
|
||||
# 如果节点名不在outputs中，创建新的节点条目
|
||||
if node_name not in spec["outputs"]:
|
||||
spec["outputs"][node_name] = {
|
||||
"node_id": node_id,
|
||||
"class_type": node.get("type"),
|
||||
}
|
||||
|
||||
if "outputs" in node:
|
||||
for an_output in node.get("outputs", []):
|
||||
output_name = an_output["name"].lower()
|
||||
|
||||
# 确保outputs字段存在
|
||||
if "outputs" not in spec["outputs"][node_name]:
|
||||
spec["outputs"][node_name]["outputs"] = {}
|
||||
|
||||
# 直接使用output_name作为字段名
|
||||
spec["outputs"][node_name]["outputs"][output_name] = {
|
||||
"output_name": an_output["name"],
|
||||
}
|
||||
|
||||
return spec
|
||||
|
||||
|
||||
def parse_inputs_json_schema(workflow_data: dict) -> dict:
|
||||
"""
|
||||
解析工作流，生成符合 JSON Schema 标准的 API 规范。
|
||||
返回一个只包含输入参数的 JSON Schema 对象。
|
||||
"""
|
||||
# 基础 JSON Schema 结构，只包含输入
|
||||
schema = {
|
||||
"$schema": "http://json-schema.org/draft-07/schema#",
|
||||
"type": "object",
|
||||
"title": "ComfyUI Workflow Input Schema",
|
||||
"description": "工作流输入参数定义",
|
||||
"properties": {},
|
||||
"required": [],
|
||||
}
|
||||
|
||||
if "nodes" not in workflow_data or not isinstance(workflow_data["nodes"], list):
|
||||
raise ValueError(
|
||||
"Invalid workflow format: 'nodes' key not found or is not a list."
|
||||
)
|
||||
|
||||
nodes_map = {str(node["id"]): node for node in workflow_data["nodes"]}
|
||||
|
||||
# 收集所有输入节点
|
||||
input_nodes = {}
|
||||
|
||||
for node_id, node in nodes_map.items():
|
||||
title: str = node.get("title")
|
||||
if not title:
|
||||
continue
|
||||
|
||||
if title.startswith(API_INPUT_PREFIX):
|
||||
# 提取节点名（去掉API_INPUT_PREFIX前缀）
|
||||
node_name = title[len(API_INPUT_PREFIX) :].lower()
|
||||
|
||||
if node_name not in input_nodes:
|
||||
input_nodes[node_name] = {
|
||||
"node_id": node_id,
|
||||
"class_type": node.get("type"),
|
||||
"properties": {},
|
||||
"required": [],
|
||||
}
|
||||
|
||||
if "inputs" in node:
|
||||
for a_input in node.get("inputs", []):
|
||||
if a_input.get("link") is None and "widget" in a_input:
|
||||
widget_name = a_input["widget"]["name"].lower()
|
||||
|
||||
# 特殊处理：隐藏LoadImage节点的upload字段
|
||||
if node.get("type") == "LoadImage" and widget_name == "upload":
|
||||
continue
|
||||
|
||||
# 生成字段的 JSON Schema 定义
|
||||
field_schema = _generate_field_schema(a_input)
|
||||
input_nodes[node_name]["properties"][widget_name] = field_schema
|
||||
|
||||
# 如果字段是必需的，添加到required列表
|
||||
if a_input.get("required", True): # 默认所有字段都是必需的
|
||||
input_nodes[node_name]["required"].append(widget_name)
|
||||
|
||||
# 将收集的节点信息转换为 JSON Schema 格式
|
||||
for node_name, node_info in input_nodes.items():
|
||||
schema["properties"][node_name] = {
|
||||
"type": "object",
|
||||
"title": f"输入节点: {node_name}",
|
||||
"description": f"节点类型: {node_info['class_type']}",
|
||||
"properties": node_info["properties"],
|
||||
"required": node_info["required"] if node_info["required"] else [],
|
||||
"additionalProperties": False,
|
||||
}
|
||||
|
||||
# 如果节点有必需字段，将整个节点标记为必需
|
||||
if node_info["required"]:
|
||||
schema["required"].append(node_name)
|
||||
|
||||
return schema
|
||||
|
||||
|
||||
########################################################
|
||||
# 以下是内部函数
|
||||
########################################################
|
||||
|
||||
|
||||
async def _patch_workflow(
|
||||
workflow_data: dict,
|
||||
api_spec: dict[str, dict[str, Any]],
|
||||
request_data: dict[str, Any],
|
||||
server: ComfyUIServer,
|
||||
) -> dict:
|
||||
"""
|
||||
将request_data中的参数值，patch到workflow_data中。并返回修改后的workflow_data。
|
||||
|
||||
request_data结构: {"节点名称": {"字段名称": "字段值"}}
|
||||
"""
|
||||
if "nodes" not in workflow_data:
|
||||
raise ValueError("无效的工作流格式")
|
||||
|
||||
nodes_map = {str(node["id"]): node for node in workflow_data["nodes"]}
|
||||
|
||||
for node_name, node_fields in request_data.items():
|
||||
if node_name not in api_spec["inputs"]:
|
||||
continue
|
||||
|
||||
node_spec = api_spec["inputs"][node_name]
|
||||
node_id = node_spec["node_id"]
|
||||
if node_id not in nodes_map:
|
||||
continue
|
||||
|
||||
target_node = nodes_map[node_id]
|
||||
|
||||
# 处理该节点下的所有字段
|
||||
for field_name, value in node_fields.items():
|
||||
if "inputs" not in node_spec or field_name not in node_spec["inputs"]:
|
||||
continue
|
||||
|
||||
field_spec = node_spec["inputs"][field_name]
|
||||
widget_name_to_patch = field_spec["widget_name"]
|
||||
|
||||
# 找到所有属于widget的输入，并确定目标widget的索引
|
||||
widget_inputs = [
|
||||
inp for inp in target_node.get("inputs", []) if "widget" in inp
|
||||
]
|
||||
|
||||
target_widget_index = -1
|
||||
for i, widget_input in enumerate(widget_inputs):
|
||||
# "widget" 字典可能不存在或没有 "name" 键
|
||||
if widget_input.get("widget", {}).get("name") == widget_name_to_patch:
|
||||
target_widget_index = i
|
||||
break
|
||||
|
||||
if target_widget_index == -1:
|
||||
logger.warning(
|
||||
f"在节点 {node_id} 中未找到名为 '{widget_name_to_patch}' 的 widget。跳过此参数。"
|
||||
)
|
||||
continue
|
||||
|
||||
# 确保 `widgets_values` 存在且为列表
|
||||
if "widgets_values" not in target_node or not isinstance(
|
||||
target_node.get("widgets_values"), list
|
||||
):
|
||||
# 如果不存在或格式错误，根据widget数量创建一个占位符列表
|
||||
target_node["widgets_values"] = [None] * len(widget_inputs)
|
||||
|
||||
# 确保 `widgets_values` 列表足够长
|
||||
while len(target_node["widgets_values"]) <= target_widget_index:
|
||||
target_node["widgets_values"].append(None)
|
||||
|
||||
# 根据API规范转换数据类型
|
||||
target_type = str
|
||||
if field_spec["type"] == "int":
|
||||
target_type = int
|
||||
elif field_spec["type"] == "float":
|
||||
target_type = float
|
||||
|
||||
# 在正确的位置上更新值
|
||||
try:
|
||||
if target_node.get("type") == "LoadImage":
|
||||
value = await _upload_image_to_comfy(value, server)
|
||||
target_node["widgets_values"][target_widget_index] = target_type(value)
|
||||
except (ValueError, TypeError) as e:
|
||||
logger.warning(
|
||||
f"无法将节点 '{node_name}' 的字段 '{field_name}' 的值 '{value}' 转换为类型 '{field_spec['type']}'。错误: {e}"
|
||||
)
|
||||
continue
|
||||
|
||||
workflow_data["nodes"] = list(nodes_map.values())
|
||||
return workflow_data
|
||||
|
||||
|
||||
def _convert_workflow_to_prompt_api_format(workflow_data: dict) -> dict:
|
||||
"""
|
||||
将工作流（API格式）转换为提交到/prompt端点的格式。
|
||||
此函数现在能正确处理已通过 `patch_workflow` 修改的 `widgets_values`。
|
||||
"""
|
||||
if "nodes" not in workflow_data:
|
||||
raise ValueError("无效的工作流格式")
|
||||
|
||||
prompt_api_format = {}
|
||||
|
||||
# 建立从link_id到源节点的映射
|
||||
link_map = {}
|
||||
for link_data in workflow_data.get("links", []):
|
||||
(
|
||||
link_id,
|
||||
origin_node_id,
|
||||
origin_slot_index,
|
||||
target_node_id,
|
||||
target_slot_index,
|
||||
link_type,
|
||||
) = link_data
|
||||
|
||||
# 键是目标节点的输入link_id
|
||||
link_map[link_id] = [str(origin_node_id), origin_slot_index]
|
||||
|
||||
for node in workflow_data["nodes"]:
|
||||
node_id = str(node["id"])
|
||||
inputs_dict = {}
|
||||
|
||||
# 1. 处理控件输入 (widgets)
|
||||
widgets_values = node.get("widgets_values", [])
|
||||
widget_cursor = 0
|
||||
for input_config in node.get("inputs", []):
|
||||
# 如果是widget并且有对应的widgets_values
|
||||
if "widget" in input_config and widget_cursor < len(widgets_values):
|
||||
widget_name = input_config["widget"].get("name")
|
||||
if widget_name:
|
||||
# 使用widgets_values中的值，因为这里已经包含了API传入的修改
|
||||
inputs_dict[widget_name] = widgets_values[widget_cursor]
|
||||
widget_cursor += 1
|
||||
|
||||
# 2. 处理连接输入 (links)
|
||||
for input_config in node.get("inputs", []):
|
||||
if "link" in input_config and input_config["link"] is not None:
|
||||
link_id = input_config["link"]
|
||||
if link_id in link_map:
|
||||
# 输入名称是input_config中的'name'
|
||||
inputs_dict[input_config["name"]] = link_map[link_id]
|
||||
|
||||
prompt_api_format[node_id] = {"class_type": node["type"], "inputs": inputs_dict}
|
||||
|
||||
return prompt_api_format
|
||||
|
||||
|
||||
async def _upload_image_to_comfy(file_path: str, server: ComfyUIServer) -> str:
|
||||
"""
|
||||
上传文件到服务器。
|
||||
file 是一个http链接
|
||||
返回一个名称
|
||||
"""
|
||||
file_name = file_path.split("/")[-1]
|
||||
file_data = await _download_file(file_path)
|
||||
|
||||
form_data = aiohttp.FormData()
|
||||
form_data.add_field("image", file_data, filename=file_name, content_type="image/*")
|
||||
form_data.add_field("type", "input")
|
||||
|
||||
async with aiohttp.ClientSession() as session:
|
||||
async with session.post(
|
||||
f"{server.http_url}/api/upload/image", data=form_data
|
||||
) as response:
|
||||
response.raise_for_status()
|
||||
result = await response.json()
|
||||
return result["name"]
|
||||
|
||||
|
||||
async def _download_file(src: str) -> bytes:
|
||||
"""
|
||||
下载文件到本地。
|
||||
"""
|
||||
async with aiohttp.ClientSession() as session:
|
||||
async with session.get(src) as response:
|
||||
response.raise_for_status()
|
||||
return await response.read()
|
||||
|
||||
|
||||
def _generate_field_schema(input_field: dict) -> dict:
|
||||
"""
|
||||
根据输入字段信息生成 JSON Schema 字段定义
|
||||
"""
|
||||
field_schema = {
|
||||
"title": input_field["widget"]["name"],
|
||||
"description": input_field.get("description", ""),
|
||||
}
|
||||
|
||||
# 根据字段类型设置 schema 类型
|
||||
field_type = _get_param_type(input_field.get("type", "STRING"))
|
||||
|
||||
if field_type == "int":
|
||||
field_schema.update(
|
||||
{
|
||||
"type": "integer",
|
||||
"minimum": input_field.get("min", None),
|
||||
"maximum": input_field.get("max", None),
|
||||
}
|
||||
)
|
||||
elif field_type == "float":
|
||||
field_schema.update(
|
||||
{
|
||||
"type": "number",
|
||||
"minimum": input_field.get("min", None),
|
||||
"maximum": input_field.get("max", None),
|
||||
}
|
||||
)
|
||||
elif field_type == "UploadFile":
|
||||
field_schema.update(
|
||||
{"type": "string", "format": "binary", "description": "上传文件"}
|
||||
)
|
||||
else: # string
|
||||
field_schema.update({"type": "string"})
|
||||
|
||||
# 处理枚举值
|
||||
if "options" in input_field:
|
||||
field_schema["enum"] = input_field["options"]
|
||||
field_schema[
|
||||
"description"
|
||||
] += f" 可选值: {', '.join(input_field['options'])}"
|
||||
|
||||
# 设置默认值
|
||||
if "default" in input_field:
|
||||
field_schema["default"] = input_field["default"]
|
||||
|
||||
return field_schema
|
||||
|
||||
|
||||
def _get_param_type(input_type_str: str) -> str:
|
||||
"""
|
||||
根据输入类型字符串确定参数类型
|
||||
"""
|
||||
input_type_str = input_type_str.upper()
|
||||
if "COMBO" in input_type_str:
|
||||
return "UploadFile"
|
||||
elif "INT" in input_type_str:
|
||||
return "int"
|
||||
elif "FLOAT" in input_type_str:
|
||||
return "float"
|
||||
else:
|
||||
return "string"
|
||||
@@ -1,855 +0,0 @@
|
||||
import asyncio
|
||||
import base64
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import random
|
||||
import uuid
|
||||
import websockets
|
||||
from datetime import datetime
|
||||
from typing import Dict, Any
|
||||
|
||||
import aiohttp
|
||||
from aiohttp import ClientTimeout
|
||||
|
||||
from workflow_service.config import Settings, ComfyUIServer
|
||||
from workflow_service.database.api import (
|
||||
create_workflow_run,
|
||||
update_workflow_run_status,
|
||||
create_workflow_run_nodes,
|
||||
update_workflow_run_node_status,
|
||||
get_workflow_run,
|
||||
get_workflow_run_nodes,
|
||||
)
|
||||
|
||||
settings = Settings()
|
||||
|
||||
API_INPUT_PREFIX = "INPUT_"
|
||||
API_OUTPUT_PREFIX = "OUTPUT_"
|
||||
|
||||
# 配置日志
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
# 全局任务队列管理器
|
||||
class WorkflowQueueManager:
|
||||
def __init__(self):
|
||||
self.running_tasks = {} # server_url -> task_info
|
||||
self.pending_tasks = [] # 等待队列
|
||||
self.lock = asyncio.Lock()
|
||||
|
||||
async def add_task(
|
||||
self,
|
||||
workflow_run_id: str,
|
||||
workflow_name: str,
|
||||
workflow_data: dict,
|
||||
api_spec: dict,
|
||||
request_data: dict,
|
||||
):
|
||||
"""添加新任务到队列"""
|
||||
async with self.lock:
|
||||
# 创建任务记录
|
||||
await create_workflow_run(
|
||||
workflow_run_id=workflow_run_id,
|
||||
workflow_name=workflow_name,
|
||||
workflow_json=json.dumps(workflow_data),
|
||||
api_spec=json.dumps(api_spec),
|
||||
request_data=json.dumps(request_data),
|
||||
)
|
||||
|
||||
# 创建工作流节点记录
|
||||
nodes_data = []
|
||||
for node in workflow_data.get("nodes", []):
|
||||
nodes_data.append(
|
||||
{"id": str(node["id"]), "type": node.get("type", "unknown")}
|
||||
)
|
||||
await create_workflow_run_nodes(workflow_run_id, nodes_data)
|
||||
|
||||
# 添加到待处理队列
|
||||
self.pending_tasks.append(workflow_run_id)
|
||||
logger.info(
|
||||
f"任务 {workflow_run_id} 已添加到队列，当前队列长度: {len(self.pending_tasks)}"
|
||||
)
|
||||
|
||||
# 尝试处理队列
|
||||
asyncio.create_task(self._process_queue())
|
||||
|
||||
return workflow_run_id
|
||||
|
||||
async def _process_queue(self):
|
||||
"""处理队列中的任务"""
|
||||
async with self.lock:
|
||||
if not self.pending_tasks:
|
||||
return
|
||||
|
||||
# 检查是否有空闲的服务器
|
||||
available_servers = await self._get_available_servers()
|
||||
if not available_servers:
|
||||
logger.info("没有可用的服务器，等待中...")
|
||||
return
|
||||
|
||||
# 获取一个待处理任务
|
||||
workflow_run_id = self.pending_tasks.pop(0)
|
||||
server = available_servers[0]
|
||||
|
||||
# 标记任务为运行中
|
||||
await update_workflow_run_status(
|
||||
workflow_run_id, "running", server.http_url
|
||||
)
|
||||
self.running_tasks[server.http_url] = {
|
||||
"workflow_run_id": workflow_run_id,
|
||||
"started_at": datetime.now(),
|
||||
}
|
||||
|
||||
# 启动任务执行
|
||||
asyncio.create_task(self._execute_task(workflow_run_id, server))
|
||||
|
||||
async def _get_available_servers(self) -> list[ComfyUIServer]:
|
||||
"""获取可用的服务器"""
|
||||
available_servers = []
|
||||
for server in settings.SERVERS:
|
||||
if server.http_url not in self.running_tasks:
|
||||
# 检查服务器状态
|
||||
try:
|
||||
async with aiohttp.ClientSession() as session:
|
||||
status = await get_server_status(server, session)
|
||||
if status["is_reachable"] and status["is_free"]:
|
||||
available_servers.append(server)
|
||||
except Exception as e:
|
||||
logger.warning(f"检查服务器 {server.http_url} 状态时出错: {e}")
|
||||
|
||||
return available_servers
|
||||
|
||||
async def _execute_task(self, workflow_run_id: str, server: ComfyUIServer):
|
||||
"""执行任务"""
|
||||
cleanup_paths = []
|
||||
try:
|
||||
# 获取工作流数据
|
||||
workflow_run = await get_workflow_run(workflow_run_id)
|
||||
if not workflow_run:
|
||||
raise Exception(f"找不到工作流运行记录: {workflow_run_id}")
|
||||
|
||||
workflow_data = json.loads(workflow_run["workflow_json"])
|
||||
api_spec = json.loads(workflow_run["api_spec"])
|
||||
request_data = json.loads(workflow_run["request_data"])
|
||||
|
||||
# 执行工作流
|
||||
result = await execute_prompt_on_server(
|
||||
workflow_data, api_spec, request_data, server, workflow_run_id
|
||||
)
|
||||
|
||||
# 保存处理后的结果到数据库
|
||||
await update_workflow_run_status(
|
||||
workflow_run_id,
|
||||
"completed",
|
||||
result=json.dumps(result, ensure_ascii=False),
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"执行任务 {workflow_run_id} 时出错: {e}")
|
||||
await update_workflow_run_status(
|
||||
workflow_run_id, "failed", error_message=str(e)
|
||||
)
|
||||
finally:
|
||||
# 清理临时文件
|
||||
if cleanup_paths:
|
||||
logger.info(f"正在清理 {len(cleanup_paths)} 个临时文件...")
|
||||
for path in cleanup_paths:
|
||||
try:
|
||||
if os.path.exists(path):
|
||||
os.remove(path)
|
||||
logger.info(f" - 已删除: {path}")
|
||||
except Exception as e:
|
||||
logger.warning(f" - 删除 {path} 时出错: {e}")
|
||||
|
||||
# 清理运行状态
|
||||
async with self.lock:
|
||||
if server.http_url in self.running_tasks:
|
||||
del self.running_tasks[server.http_url]
|
||||
|
||||
# 继续处理队列
|
||||
asyncio.create_task(self._process_queue())
|
||||
|
||||
async def get_task_status(self, workflow_run_id: str) -> dict:
|
||||
"""获取任务状态"""
|
||||
workflow_run = await get_workflow_run(workflow_run_id)
|
||||
if not workflow_run:
|
||||
return {"error": "任务不存在"}
|
||||
nodes = await get_workflow_run_nodes(workflow_run_id)
|
||||
|
||||
result = {
|
||||
"id": workflow_run_id,
|
||||
"status": workflow_run["status"],
|
||||
"created_at": workflow_run["created_at"],
|
||||
"started_at": workflow_run["started_at"],
|
||||
"completed_at": workflow_run["completed_at"],
|
||||
"server_url": workflow_run["server_url"],
|
||||
"error_message": workflow_run["error_message"],
|
||||
"nodes": nodes,
|
||||
}
|
||||
|
||||
# 如果任务完成，从数据库获取结果
|
||||
if workflow_run["status"] == "completed" and workflow_run.get("result"):
|
||||
try:
|
||||
result["result"] = json.loads(workflow_run["result"])
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
result["result"] = None
|
||||
|
||||
return result
|
||||
|
||||
|
||||
# 全局队列管理器实例
|
||||
queue_manager = WorkflowQueueManager()
|
||||
|
||||
|
||||
# 定义一个自定义异常，用于封装来自ComfyUI的执行错误
|
||||
class ComfyUIExecutionError(Exception):
|
||||
def __init__(self, error_data: dict):
|
||||
self.error_data = error_data
|
||||
# 创建一个对开发者友好的异常消息
|
||||
message = (
|
||||
f"ComfyUI节点执行失败。节点ID: {error_data.get('node_id')}, "
|
||||
f"节点类型: {error_data.get('node_type')}. "
|
||||
f"错误: {error_data.get('exception_message', 'N/A')}"
|
||||
)
|
||||
super().__init__(message)
|
||||
|
||||
|
||||
async def get_server_status(
|
||||
server: ComfyUIServer, session: aiohttp.ClientSession
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
检查单个ComfyUI服务器的详细状态。
|
||||
返回一个包含可达性、队列状态和详细队列内容的字典。
|
||||
"""
|
||||
# [BUG修复] 确保初始字典结构与成功时的结构一致，以满足Pydantic模型
|
||||
status_info = {
|
||||
"is_reachable": False,
|
||||
"is_free": False,
|
||||
"queue_details": {"running_count": 0, "pending_count": 0},
|
||||
}
|
||||
try:
|
||||
queue_url = f"{server.http_url}/queue"
|
||||
async with session.get(queue_url, timeout=60) as response:
|
||||
response.raise_for_status()
|
||||
queue_data = await response.json()
|
||||
|
||||
status_info["is_reachable"] = True
|
||||
|
||||
running_count = len(queue_data.get("queue_running", []))
|
||||
pending_count = len(queue_data.get("queue_pending", []))
|
||||
|
||||
status_info["queue_details"] = {
|
||||
"running_count": running_count,
|
||||
"pending_count": pending_count,
|
||||
}
|
||||
|
||||
status_info["is_free"] = running_count == 0 and pending_count == 0
|
||||
|
||||
except Exception as e:
|
||||
# 当请求失败时，将返回上面定义的、结构正确的初始 status_info
|
||||
logger.warning(f"无法检查服务器 {server.http_url} 的队列状态: {e}")
|
||||
|
||||
return status_info
|
||||
|
||||
|
||||
async def execute_prompt_on_server(
|
||||
workflow_data: Dict,
|
||||
api_spec: Dict,
|
||||
request_data: Dict,
|
||||
server: ComfyUIServer,
|
||||
workflow_run_id: str,
|
||||
) -> Dict:
|
||||
"""
|
||||
在指定的服务器上执行一个准备好的prompt。
|
||||
现在支持节点级别的状态跟踪。
|
||||
"""
|
||||
client_id = str(uuid.uuid4())
|
||||
|
||||
# 应用请求数据到工作流
|
||||
patched_workflow = await patch_workflow(
|
||||
workflow_data, api_spec, request_data, server
|
||||
)
|
||||
|
||||
# 转换为prompt格式
|
||||
prompt = convert_workflow_to_prompt_api_format(patched_workflow)
|
||||
|
||||
# 更新工作流运行状态，记录prompt_id和client_id
|
||||
await update_workflow_run_status(
|
||||
workflow_run_id,
|
||||
"running",
|
||||
server.http_url,
|
||||
None, # prompt_id将在_queue_prompt后更新
|
||||
client_id,
|
||||
)
|
||||
|
||||
# 提交到ComfyUI
|
||||
prompt_id = await _queue_prompt(workflow_data, prompt, client_id, server.http_url)
|
||||
|
||||
# 更新prompt_id
|
||||
await update_workflow_run_status(
|
||||
workflow_run_id, "running", server.http_url, prompt_id, client_id
|
||||
)
|
||||
|
||||
logger.info(
|
||||
f"工作流 {workflow_run_id} 已在 {server.http_url} 上入队，Prompt ID: {prompt_id}"
|
||||
)
|
||||
|
||||
# 获取执行结果，现在支持节点级别的状态跟踪
|
||||
results = await _get_execution_results(
|
||||
workflow_data, prompt_id, client_id, server.ws_url, workflow_run_id
|
||||
)
|
||||
return results
|
||||
|
||||
|
||||
async def submit_workflow_to_queue(
|
||||
workflow_name: str, workflow_data: Dict, api_spec: Dict, request_data: Dict
|
||||
) -> str:
|
||||
"""
|
||||
提交工作流到队列，立即返回任务ID。
|
||||
这是新的异步接口，调用者可以通过任务ID查询状态。
|
||||
"""
|
||||
workflow_run_id = str(uuid.uuid4())
|
||||
|
||||
# 添加到队列管理器
|
||||
await queue_manager.add_task(
|
||||
workflow_run_id, workflow_name, workflow_data, api_spec, request_data
|
||||
)
|
||||
|
||||
return workflow_run_id
|
||||
|
||||
|
||||
def parse_api_spec(workflow_data: dict) -> Dict[str, Dict[str, Any]]:
|
||||
"""
|
||||
解析工作流，并根据规范生成嵌套结构的API参数名。
|
||||
结构: {'节点名': {'node_id': '节点ID', '字段名': {字段信息}}}
|
||||
"""
|
||||
spec = {"inputs": {}, "outputs": {}}
|
||||
if "nodes" not in workflow_data or not isinstance(workflow_data["nodes"], list):
|
||||
raise ValueError(
|
||||
"Invalid workflow format: 'nodes' key not found or is not a list."
|
||||
)
|
||||
|
||||
nodes_map = {str(node["id"]): node for node in workflow_data["nodes"]}
|
||||
|
||||
for node_id, node in nodes_map.items():
|
||||
title: str = node.get("title")
|
||||
if not title:
|
||||
continue
|
||||
|
||||
if title.startswith(API_INPUT_PREFIX):
|
||||
# 提取节点名（去掉API_INPUT_PREFIX前缀）
|
||||
node_name = title[len(API_INPUT_PREFIX) :].lower()
|
||||
|
||||
# 如果节点名不在inputs中，创建新的节点条目
|
||||
if node_name not in spec["inputs"]:
|
||||
spec["inputs"][node_name] = {
|
||||
"node_id": node_id,
|
||||
"class_type": node.get("type"),
|
||||
}
|
||||
|
||||
if "inputs" in node:
|
||||
for a_input in node.get("inputs", []):
|
||||
if a_input.get("link") is None and "widget" in a_input:
|
||||
widget_name = a_input["widget"]["name"].lower()
|
||||
|
||||
# 特殊处理：隐藏LoadImage节点的upload字段
|
||||
if node.get("type") == "LoadImage" and widget_name == "upload":
|
||||
continue
|
||||
|
||||
# 确保inputs字段存在
|
||||
if "inputs" not in spec["inputs"][node_name]:
|
||||
spec["inputs"][node_name]["inputs"] = {}
|
||||
|
||||
# 直接使用widget_name作为字段名
|
||||
spec["inputs"][node_name]["inputs"][widget_name] = {
|
||||
"type": _get_param_type(a_input.get("type", "STRING")),
|
||||
"widget_name": a_input["widget"]["name"],
|
||||
}
|
||||
|
||||
elif title.startswith(API_OUTPUT_PREFIX):
|
||||
# 提取节点名（去掉API_OUTPUT_PREFIX前缀）
|
||||
node_name = title[len(API_OUTPUT_PREFIX) :].lower()
|
||||
|
||||
# 如果节点名不在outputs中，创建新的节点条目
|
||||
if node_name not in spec["outputs"]:
|
||||
spec["outputs"][node_name] = {
|
||||
"node_id": node_id,
|
||||
"class_type": node.get("type"),
|
||||
}
|
||||
|
||||
if "outputs" in node:
|
||||
for an_output in node.get("outputs", []):
|
||||
output_name = an_output["name"].lower()
|
||||
|
||||
# 确保outputs字段存在
|
||||
if "outputs" not in spec["outputs"][node_name]:
|
||||
spec["outputs"][node_name]["outputs"] = {}
|
||||
|
||||
# 直接使用output_name作为字段名
|
||||
spec["outputs"][node_name]["outputs"][output_name] = {
|
||||
"output_name": an_output["name"],
|
||||
}
|
||||
|
||||
return spec
|
||||
|
||||
|
||||
def parse_inputs_json_schema(workflow_data: dict) -> dict:
|
||||
"""
|
||||
解析工作流，生成符合 JSON Schema 标准的 API 规范。
|
||||
返回一个只包含输入参数的 JSON Schema 对象。
|
||||
"""
|
||||
# 基础 JSON Schema 结构，只包含输入
|
||||
schema = {
|
||||
"$schema": "http://json-schema.org/draft-07/schema#",
|
||||
"type": "object",
|
||||
"title": "ComfyUI Workflow Input Schema",
|
||||
"description": "工作流输入参数定义",
|
||||
"properties": {},
|
||||
"required": [],
|
||||
}
|
||||
|
||||
if "nodes" not in workflow_data or not isinstance(workflow_data["nodes"], list):
|
||||
raise ValueError(
|
||||
"Invalid workflow format: 'nodes' key not found or is not a list."
|
||||
)
|
||||
|
||||
nodes_map = {str(node["id"]): node for node in workflow_data["nodes"]}
|
||||
|
||||
# 收集所有输入节点
|
||||
input_nodes = {}
|
||||
|
||||
for node_id, node in nodes_map.items():
|
||||
title: str = node.get("title")
|
||||
if not title:
|
||||
continue
|
||||
|
||||
if title.startswith(API_INPUT_PREFIX):
|
||||
# 提取节点名（去掉API_INPUT_PREFIX前缀）
|
||||
node_name = title[len(API_INPUT_PREFIX) :].lower()
|
||||
|
||||
if node_name not in input_nodes:
|
||||
input_nodes[node_name] = {
|
||||
"node_id": node_id,
|
||||
"class_type": node.get("type"),
|
||||
"properties": {},
|
||||
"required": [],
|
||||
}
|
||||
|
||||
if "inputs" in node:
|
||||
for a_input in node.get("inputs", []):
|
||||
if a_input.get("link") is None and "widget" in a_input:
|
||||
widget_name = a_input["widget"]["name"].lower()
|
||||
|
||||
# 特殊处理：隐藏LoadImage节点的upload字段
|
||||
if node.get("type") == "LoadImage" and widget_name == "upload":
|
||||
continue
|
||||
|
||||
# 生成字段的 JSON Schema 定义
|
||||
field_schema = _generate_field_schema(a_input)
|
||||
input_nodes[node_name]["properties"][widget_name] = field_schema
|
||||
|
||||
# 如果字段是必需的，添加到required列表
|
||||
if a_input.get("required", True): # 默认所有字段都是必需的
|
||||
input_nodes[node_name]["required"].append(widget_name)
|
||||
|
||||
# 将收集的节点信息转换为 JSON Schema 格式
|
||||
for node_name, node_info in input_nodes.items():
|
||||
schema["properties"][node_name] = {
|
||||
"type": "object",
|
||||
"title": f"输入节点: {node_name}",
|
||||
"description": f"节点类型: {node_info['class_type']}",
|
||||
"properties": node_info["properties"],
|
||||
"required": node_info["required"] if node_info["required"] else [],
|
||||
"additionalProperties": False,
|
||||
}
|
||||
|
||||
# 如果节点有必需字段，将整个节点标记为必需
|
||||
if node_info["required"]:
|
||||
schema["required"].append(node_name)
|
||||
|
||||
return schema
|
||||
|
||||
|
||||
def _generate_field_schema(input_field: dict) -> dict:
|
||||
"""
|
||||
根据输入字段信息生成 JSON Schema 字段定义
|
||||
"""
|
||||
field_schema = {
|
||||
"title": input_field["widget"]["name"],
|
||||
"description": input_field.get("description", ""),
|
||||
}
|
||||
|
||||
# 根据字段类型设置 schema 类型
|
||||
field_type = _get_param_type(input_field.get("type", "STRING"))
|
||||
|
||||
if field_type == "int":
|
||||
field_schema.update(
|
||||
{
|
||||
"type": "integer",
|
||||
"minimum": input_field.get("min", None),
|
||||
"maximum": input_field.get("max", None),
|
||||
}
|
||||
)
|
||||
elif field_type == "float":
|
||||
field_schema.update(
|
||||
{
|
||||
"type": "number",
|
||||
"minimum": input_field.get("min", None),
|
||||
"maximum": input_field.get("max", None),
|
||||
}
|
||||
)
|
||||
elif field_type == "UploadFile":
|
||||
field_schema.update(
|
||||
{"type": "string", "format": "binary", "description": "上传文件"}
|
||||
)
|
||||
else: # string
|
||||
field_schema.update({"type": "string"})
|
||||
|
||||
# 处理枚举值
|
||||
if "options" in input_field:
|
||||
field_schema["enum"] = input_field["options"]
|
||||
field_schema[
|
||||
"description"
|
||||
] += f" 可选值: {', '.join(input_field['options'])}"
|
||||
|
||||
# 设置默认值
|
||||
if "default" in input_field:
|
||||
field_schema["default"] = input_field["default"]
|
||||
|
||||
return field_schema
|
||||
|
||||
|
||||
def _get_param_type(input_type_str: str) -> str:
|
||||
"""
|
||||
根据输入类型字符串确定参数类型
|
||||
"""
|
||||
input_type_str = input_type_str.upper()
|
||||
if "COMBO" in input_type_str:
|
||||
return "UploadFile"
|
||||
elif "INT" in input_type_str:
|
||||
return "int"
|
||||
elif "FLOAT" in input_type_str:
|
||||
return "float"
|
||||
else:
|
||||
return "string"
|
||||
|
||||
|
||||
async def patch_workflow(
|
||||
workflow_data: dict,
|
||||
api_spec: dict[str, dict[str, Any]],
|
||||
request_data: dict[str, Any],
|
||||
server: ComfyUIServer,
|
||||
) -> dict:
|
||||
"""
|
||||
将request_data中的参数值，patch到workflow_data中。并返回修改后的workflow_data。
|
||||
|
||||
request_data结构: {"节点名称": {"字段名称": "字段值"}}
|
||||
"""
|
||||
if "nodes" not in workflow_data:
|
||||
raise ValueError("无效的工作流格式")
|
||||
|
||||
nodes_map = {str(node["id"]): node for node in workflow_data["nodes"]}
|
||||
|
||||
for node_name, node_fields in request_data.items():
|
||||
if node_name not in api_spec["inputs"]:
|
||||
continue
|
||||
|
||||
node_spec = api_spec["inputs"][node_name]
|
||||
node_id = node_spec["node_id"]
|
||||
if node_id not in nodes_map:
|
||||
continue
|
||||
|
||||
target_node = nodes_map[node_id]
|
||||
|
||||
# 处理该节点下的所有字段
|
||||
for field_name, value in node_fields.items():
|
||||
if "inputs" not in node_spec or field_name not in node_spec["inputs"]:
|
||||
continue
|
||||
|
||||
field_spec = node_spec["inputs"][field_name]
|
||||
widget_name_to_patch = field_spec["widget_name"]
|
||||
|
||||
# 找到所有属于widget的输入，并确定目标widget的索引
|
||||
widget_inputs = [
|
||||
inp for inp in target_node.get("inputs", []) if "widget" in inp
|
||||
]
|
||||
|
||||
target_widget_index = -1
|
||||
for i, widget_input in enumerate(widget_inputs):
|
||||
# "widget" 字典可能不存在或没有 "name" 键
|
||||
if widget_input.get("widget", {}).get("name") == widget_name_to_patch:
|
||||
target_widget_index = i
|
||||
break
|
||||
|
||||
if target_widget_index == -1:
|
||||
logger.warning(
|
||||
f"在节点 {node_id} 中未找到名为 '{widget_name_to_patch}' 的 widget。跳过此参数。"
|
||||
)
|
||||
continue
|
||||
|
||||
# 确保 `widgets_values` 存在且为列表
|
||||
if "widgets_values" not in target_node or not isinstance(
|
||||
target_node.get("widgets_values"), list
|
||||
):
|
||||
# 如果不存在或格式错误，根据widget数量创建一个占位符列表
|
||||
target_node["widgets_values"] = [None] * len(widget_inputs)
|
||||
|
||||
# 确保 `widgets_values` 列表足够长
|
||||
while len(target_node["widgets_values"]) <= target_widget_index:
|
||||
target_node["widgets_values"].append(None)
|
||||
|
||||
# 根据API规范转换数据类型
|
||||
target_type = str
|
||||
if field_spec["type"] == "int":
|
||||
target_type = int
|
||||
elif field_spec["type"] == "float":
|
||||
target_type = float
|
||||
|
||||
# 在正确的位置上更新值
|
||||
try:
|
||||
if target_node.get("type") == "LoadImage":
|
||||
value = await upload_image_to_comfy(value, server)
|
||||
target_node["widgets_values"][target_widget_index] = target_type(value)
|
||||
except (ValueError, TypeError) as e:
|
||||
logger.warning(
|
||||
f"无法将节点 '{node_name}' 的字段 '{field_name}' 的值 '{value}' 转换为类型 '{field_spec['type']}'。错误: {e}"
|
||||
)
|
||||
continue
|
||||
|
||||
workflow_data["nodes"] = list(nodes_map.values())
|
||||
return workflow_data
|
||||
|
||||
|
||||
def convert_workflow_to_prompt_api_format(workflow_data: dict) -> dict:
|
||||
"""
|
||||
将工作流（API格式）转换为提交到/prompt端点的格式。
|
||||
此函数现在能正确处理已通过 `patch_workflow` 修改的 `widgets_values`。
|
||||
"""
|
||||
if "nodes" not in workflow_data:
|
||||
raise ValueError("无效的工作流格式")
|
||||
|
||||
prompt_api_format = {}
|
||||
|
||||
# 建立从link_id到源节点的映射
|
||||
link_map = {}
|
||||
for link_data in workflow_data.get("links", []):
|
||||
(
|
||||
link_id,
|
||||
origin_node_id,
|
||||
origin_slot_index,
|
||||
target_node_id,
|
||||
target_slot_index,
|
||||
link_type,
|
||||
) = link_data
|
||||
|
||||
# 键是目标节点的输入link_id
|
||||
link_map[link_id] = [str(origin_node_id), origin_slot_index]
|
||||
|
||||
for node in workflow_data["nodes"]:
|
||||
node_id = str(node["id"])
|
||||
inputs_dict = {}
|
||||
|
||||
# 1. 处理控件输入 (widgets)
|
||||
widgets_values = node.get("widgets_values", [])
|
||||
widget_cursor = 0
|
||||
for input_config in node.get("inputs", []):
|
||||
# 如果是widget并且有对应的widgets_values
|
||||
if "widget" in input_config and widget_cursor < len(widgets_values):
|
||||
widget_name = input_config["widget"].get("name")
|
||||
if widget_name:
|
||||
# 使用widgets_values中的值，因为这里已经包含了API传入的修改
|
||||
inputs_dict[widget_name] = widgets_values[widget_cursor]
|
||||
widget_cursor += 1
|
||||
|
||||
# 2. 处理连接输入 (links)
|
||||
for input_config in node.get("inputs", []):
|
||||
if "link" in input_config and input_config["link"] is not None:
|
||||
link_id = input_config["link"]
|
||||
if link_id in link_map:
|
||||
# 输入名称是input_config中的'name'
|
||||
inputs_dict[input_config["name"]] = link_map[link_id]
|
||||
|
||||
prompt_api_format[node_id] = {"class_type": node["type"], "inputs": inputs_dict}
|
||||
|
||||
return prompt_api_format
|
||||
|
||||
|
||||
async def _queue_prompt(
|
||||
workflow: dict, prompt: dict, client_id: str, http_url: str
|
||||
) -> str:
|
||||
"""通过HTTP POST将工作流任务提交到指定的ComfyUI服务器。"""
|
||||
for node_id in prompt:
|
||||
prompt[node_id]["inputs"][f"cache_buster_{uuid.uuid4().hex}"] = random.random()
|
||||
payload = {
|
||||
"prompt": prompt,
|
||||
"client_id": client_id,
|
||||
"extra_data": {
|
||||
"api_key_comfy_org": "",
|
||||
"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:
|
||||
async with session.post(prompt_url, json=payload) as response:
|
||||
logger.info(f"ComfyUI /prompt 端点返回的响应: {response}")
|
||||
response.raise_for_status()
|
||||
result = await response.json()
|
||||
logger.info(f"ComfyUI /prompt 端点返回的响应: {result}")
|
||||
if "prompt_id" not in result:
|
||||
raise Exception(f"从 ComfyUI /prompt 端点返回的响应无效: {result}")
|
||||
return result["prompt_id"]
|
||||
except Exception as e:
|
||||
logger.error(f"提交到 ComfyUI /prompt 端点时发生错误: {e}")
|
||||
raise e
|
||||
|
||||
|
||||
async def _get_execution_results(
|
||||
workflow_data: Dict,
|
||||
prompt_id: str,
|
||||
client_id: str,
|
||||
ws_url: str,
|
||||
workflow_run_id: str,
|
||||
) -> dict:
|
||||
"""
|
||||
通过WebSocket连接到指定的ComfyUI服务器，聚合执行结果。
|
||||
现在支持节点级别的状态跟踪。
|
||||
"""
|
||||
full_ws_url = f"{ws_url}?clientId={client_id}"
|
||||
aggregated_outputs = {}
|
||||
|
||||
try:
|
||||
async with websockets.connect(full_ws_url) as websocket:
|
||||
while True:
|
||||
try:
|
||||
out = await websocket.recv()
|
||||
if not isinstance(out, str):
|
||||
continue
|
||||
|
||||
message = json.loads(out)
|
||||
msg_type = message.get("type")
|
||||
data = message.get("data")
|
||||
|
||||
if not (data and data.get("prompt_id") == prompt_id):
|
||||
continue
|
||||
|
||||
# 捕获并处理执行错误
|
||||
if msg_type == "execution_error":
|
||||
error_data = data
|
||||
logger.error(
|
||||
f"ComfyUI执行错误 (Prompt ID: {prompt_id}): {error_data}"
|
||||
)
|
||||
|
||||
# 更新节点状态为失败
|
||||
node_id = error_data.get("node_id")
|
||||
if node_id:
|
||||
await update_workflow_run_node_status(
|
||||
workflow_run_id,
|
||||
node_id,
|
||||
"failed",
|
||||
error_message=error_data.get(
|
||||
"exception_message", "Unknown error"
|
||||
),
|
||||
)
|
||||
|
||||
# 抛出自定义异常，将错误详情传递出去
|
||||
raise ComfyUIExecutionError(error_data)
|
||||
|
||||
# 处理节点开始执行
|
||||
if msg_type == "executing" and data.get("node"):
|
||||
node_id = data.get("node")
|
||||
logger.info(f"节点 {node_id} 开始执行 (Prompt ID: {prompt_id})")
|
||||
|
||||
# 更新节点状态为运行中
|
||||
await update_workflow_run_node_status(
|
||||
workflow_run_id, node_id, "running"
|
||||
)
|
||||
|
||||
# 处理节点执行完成
|
||||
if msg_type == "executed":
|
||||
node_id = data.get("node")
|
||||
output_data = data.get("output")
|
||||
if node_id and output_data:
|
||||
node = next(
|
||||
(
|
||||
x
|
||||
for x in workflow_data["nodes"]
|
||||
if str(x["id"]) == node_id
|
||||
),
|
||||
None,
|
||||
)
|
||||
if (
|
||||
node
|
||||
and node.get("title", "")
|
||||
and node["title"].startswith(API_OUTPUT_PREFIX)
|
||||
):
|
||||
aggregated_outputs[node["title"]] = output_data
|
||||
logger.info(
|
||||
f"收到节点 {node_id} 的输出 (Prompt ID: {prompt_id})"
|
||||
)
|
||||
|
||||
# 更新节点状态为完成
|
||||
await update_workflow_run_node_status(
|
||||
workflow_run_id,
|
||||
node_id,
|
||||
"completed",
|
||||
output_data=json.dumps(output_data),
|
||||
)
|
||||
|
||||
# 处理整个工作流执行完成
|
||||
elif msg_type == "executing" and data.get("node") is None:
|
||||
logger.info(f"Prompt ID: {prompt_id} 执行完成。")
|
||||
return aggregated_outputs
|
||||
|
||||
except websockets.exceptions.ConnectionClosed as e:
|
||||
logger.warning(
|
||||
f"WebSocket 连接已关闭 (Prompt ID: {prompt_id})。错误: {e}"
|
||||
)
|
||||
return aggregated_outputs
|
||||
except Exception as e:
|
||||
# 重新抛出我们自己的异常，或者处理其他意外错误
|
||||
if not isinstance(e, ComfyUIExecutionError):
|
||||
logger.error(f"处理 prompt {prompt_id} 时发生意外错误: {e}")
|
||||
raise e
|
||||
|
||||
except websockets.exceptions.InvalidURI as e:
|
||||
logger.error(
|
||||
f"错误: 尝试连接的WebSocket URI无效: '{full_ws_url}'. 原始URL: '{ws_url}'. 错误: {e}"
|
||||
)
|
||||
raise e
|
||||
return aggregated_outputs
|
||||
|
||||
|
||||
async def upload_image_to_comfy(file_path: str, server: ComfyUIServer) -> str:
|
||||
"""
|
||||
上传文件到服务器。
|
||||
file 是一个http链接
|
||||
返回一个名称
|
||||
"""
|
||||
file_name = file_path.split("/")[-1]
|
||||
file_data = await download_file(file_path)
|
||||
|
||||
form_data = aiohttp.FormData()
|
||||
form_data.add_field("image", file_data, filename=file_name, content_type="image/*")
|
||||
form_data.add_field("type", "input")
|
||||
|
||||
async with aiohttp.ClientSession() as session:
|
||||
async with session.post(
|
||||
f"{server.http_url}/api/upload/image", data=form_data
|
||||
) as response:
|
||||
response.raise_for_status()
|
||||
result = await response.json()
|
||||
return result["name"]
|
||||
|
||||
|
||||
async def download_file(src: str) -> bytes:
|
||||
"""
|
||||
下载文件到本地。
|
||||
"""
|
||||
async with aiohttp.ClientSession() as session:
|
||||
async with session.get(src) as response:
|
||||
response.raise_for_status()
|
||||
return await response.read()
|
||||
@@ -6,6 +6,7 @@ from typing import List
|
||||
|
||||
class ComfyUIServer(BaseModel):
|
||||
"""定义单个ComfyUI服务器的配置模型。"""
|
||||
|
||||
http_url: str
|
||||
ws_url: str
|
||||
input_dir: str
|
||||
@@ -16,10 +17,13 @@ class Settings(BaseSettings):
|
||||
"""
|
||||
应用配置类，通过.env文件加载环境变量。
|
||||
"""
|
||||
model_config = SettingsConfigDict(env_file='.env', env_file_encoding='utf-8')
|
||||
|
||||
model_config = SettingsConfigDict(env_file=".env", env_file_encoding="utf-8")
|
||||
|
||||
# [核心改动] 使用JSON字符串来定义一个服务器列表
|
||||
COMFYUI_SERVERS_JSON: str = '[{"http_url": "http://127.0.0.1:8188", "ws_url": "ws://127.0.0.1:8188/ws", "input_dir": "/comfyui/input", "output_dir": "/comfyui/output"}]'
|
||||
COMFYUI_SERVERS_JSON: str = (
|
||||
'[{"http_url": "http://127.0.0.1:8188", "ws_url": "ws://127.0.0.1:8188/ws", "input_dir": "/comfyui/input", "output_dir": "/comfyui/output"}]'
|
||||
)
|
||||
|
||||
DB_FILE: str = "/db/workflows_service.sqlite"
|
||||
# S3 和 AWS 配置保持不变
|
||||
@@ -45,4 +49,4 @@ class Settings(BaseSettings):
|
||||
raise e
|
||||
|
||||
|
||||
settings = Settings()
|
||||
settings = Settings()
|
||||
|
||||
@@ -4,7 +4,8 @@ from typing import Dict, Optional
|
||||
from fastapi import APIRouter, Body, HTTPException
|
||||
from fastapi.responses import JSONResponse
|
||||
|
||||
from workflow_service import comfyui_client
|
||||
from workflow_service.comfy import comfy_workflow
|
||||
from workflow_service.comfy.comfy_queue import queue_manager
|
||||
from workflow_service.database.api import get_workflow
|
||||
|
||||
run_router = APIRouter(
|
||||
@@ -38,10 +39,10 @@ async def run_workflow(
|
||||
raise HTTPException(status_code=404, detail=detail)
|
||||
|
||||
workflow = json.loads(workflow_data["workflow_json"])
|
||||
api_spec = comfyui_client.parse_api_spec(workflow)
|
||||
api_spec = comfy_workflow.parse_api_spec(workflow)
|
||||
|
||||
# 提交到队列
|
||||
workflow_run_id = await comfyui_client.submit_workflow_to_queue(
|
||||
workflow_run_id = await queue_manager.add_task(
|
||||
workflow_name=workflow_name,
|
||||
workflow_data=workflow,
|
||||
api_spec=api_spec,
|
||||
@@ -67,7 +68,7 @@ async def get_run_status(workflow_run_id: str):
|
||||
获取工作流执行状态。
|
||||
"""
|
||||
try:
|
||||
status = await comfyui_client.queue_manager.get_task_status(workflow_run_id)
|
||||
status = await queue_manager.get_task_status(workflow_run_id)
|
||||
return status
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=500, detail=f"获取状态失败: {str(e)}")
|
||||
|
||||
@@ -3,7 +3,8 @@ from typing import Optional
|
||||
|
||||
from fastapi import Body, HTTPException
|
||||
from fastapi.responses import JSONResponse
|
||||
from workflow_service import comfyui_client
|
||||
from workflow_service.comfy import comfy_workflow
|
||||
from workflow_service.comfy.comfy_queue import queue_manager
|
||||
from workflow_service.database.api import get_workflow
|
||||
from ._base import runx_router
|
||||
|
||||
@@ -32,7 +33,7 @@ async def model_with_multi_dress(
|
||||
raise HTTPException(status_code=404, detail=detail)
|
||||
|
||||
workflow = json.loads(workflow_data["workflow_json"])
|
||||
api_spec = comfyui_client.parse_api_spec(workflow)
|
||||
api_spec = comfy_workflow.parse_api_spec(workflow)
|
||||
|
||||
# 将请求拆分为多个请求
|
||||
batch_data = _convert(data)
|
||||
@@ -40,7 +41,7 @@ async def model_with_multi_dress(
|
||||
# 提交到队列
|
||||
workflow_run_ids: list[str] = []
|
||||
for item in batch_data:
|
||||
workflow_run_id = await comfyui_client.submit_workflow_to_queue(
|
||||
workflow_run_id = await queue_manager.add_task(
|
||||
workflow_name=workflow_name,
|
||||
workflow_data=workflow,
|
||||
api_spec=api_spec,
|
||||
|
||||
@@ -7,7 +7,7 @@ import aiohttp
|
||||
from fastapi import APIRouter, HTTPException, Path
|
||||
from pydantic import BaseModel
|
||||
|
||||
from workflow_service import comfyui_client
|
||||
from workflow_service.comfy.comfy_queue import queue_manager
|
||||
from workflow_service.config import Settings
|
||||
|
||||
settings = Settings()
|
||||
@@ -54,13 +54,13 @@ async def get_metrics():
|
||||
获取队列状态概览。
|
||||
"""
|
||||
try:
|
||||
pending_count = len(comfyui_client.queue_manager.pending_tasks)
|
||||
running_count = len(comfyui_client.queue_manager.running_tasks)
|
||||
pending_count = len(queue_manager.pending_tasks)
|
||||
running_count = len(queue_manager.running_tasks)
|
||||
|
||||
return {
|
||||
"pending_tasks": pending_count,
|
||||
"running_tasks": running_count,
|
||||
"total_servers": len(comfyui_client.queue_manager.running_tasks),
|
||||
"total_servers": len(queue_manager.running_tasks),
|
||||
"queue_manager_status": "active",
|
||||
}
|
||||
except Exception as e:
|
||||
@@ -78,7 +78,7 @@ async def get_servers_status():
|
||||
|
||||
async with aiohttp.ClientSession() as session:
|
||||
status_tasks = [
|
||||
comfyui_client.get_server_status(server, session) for server in servers
|
||||
queue_manager.get_server_status(server, session) for server in servers
|
||||
]
|
||||
live_statuses = await asyncio.gather(*status_tasks)
|
||||
|
||||
|
||||
@@ -4,7 +4,7 @@ from typing import Optional, List
|
||||
from fastapi import APIRouter, Body, HTTPException, Path
|
||||
from fastapi.responses import JSONResponse
|
||||
|
||||
from workflow_service import comfyui_client
|
||||
from workflow_service.comfy import comfy_workflow
|
||||
from workflow_service.database.api import (
|
||||
save_workflow,
|
||||
get_all_workflows,
|
||||
@@ -94,8 +94,8 @@ async def get_one_workflow_endpoint(base_name: str, version: Optional[str] = Non
|
||||
workflow = json.loads(workflow_data["workflow_json"])
|
||||
return {
|
||||
"workflow": workflow,
|
||||
"api_spec": comfyui_client.parse_api_spec(workflow),
|
||||
"inputs_json_schema": comfyui_client.parse_inputs_json_schema(workflow),
|
||||
"api_spec": comfy_workflow.parse_api_spec(workflow),
|
||||
"inputs_json_schema": comfy_workflow.parse_inputs_json_schema(workflow),
|
||||
}
|
||||
except Exception as e:
|
||||
raise HTTPException(
|
||||
|
||||
Reference in New Issue
Block a user