comi-logic/ComfyUI-WorkflowPublisher
7
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix

Browse Source
This commit is contained in:
iHeyTang
2025-08-13 15:16:02 +08:00
parent d9ee791922
commit 1af6b42573
5 changed files with 919 additions and 248 deletions
Download Patch File Download Diff File Expand all files Collapse all files

51
.env
View File

@@ -1,29 +1,28 @@
COMFYUI_SERVERS_JSON='[
{
"http_url": "https://bowongai-dev--comfyui-for-waas-ui-1.modal.run",
"ws_url": "wss://bowongai-dev--comfyui-for-waas-ui-1.modal.run/ws",
"input_dir": "waas/comfyui_1/input",
"output_dir": "waas/comfyui_1/output"
},
{
"http_url": "https://bowongai-dev--comfyui-for-waas-ui-2.modal.run",
"ws_url": "wss://bowongai-dev--comfyui-for-waas-ui-2.modal.run/ws",
"input_dir": "waas/comfyui_2/input",
"output_dir": "waas/comfyui_2/output"
},
{
"http_url": "https://bowongai-dev--comfyui-for-waas-ui-3.modal.run",
"ws_url": "wss://bowongai-dev--comfyui-for-waas-ui-3.modal.run/ws",
"input_dir": "waas/comfyui_3/input",
"output_dir": "waas/comfyui_3/output"
},
{
"http_url": "https://bowongai-dev--comfyui-for-waas-ui-4.modal.run",
"ws_url": "wss://bowongai-dev--comfyui-for-waas-ui-4.modal.run/ws",
"input_dir": "waas/comfyui_4/input",
"output_dir": "waas/comfyui_4/output"
}
]'
# {
# "http_url": "https://bowongai-dev--comfyui-for-waas-ui-1.modal.run",
# "ws_url": "wss://bowongai-dev--comfyui-for-waas-ui-1.modal.run/ws",
# "input_dir": "waas/comfyui_1/input",
# "output_dir": "waas/comfyui_1/output"
# },
# {
# "http_url": "https://bowongai-dev--comfyui-for-waas-ui-2.modal.run",
# "ws_url": "wss://bowongai-dev--comfyui-for-waas-ui-2.modal.run/ws",
# "input_dir": "waas/comfyui_2/input",
# "output_dir": "waas/comfyui_2/output"
# },
# {
# "http_url": "https://bowongai-dev--comfyui-for-waas-ui-3.modal.run",
# "ws_url": "wss://bowongai-dev--comfyui-for-waas-ui-3.modal.run/ws",
# "input_dir": "waas/comfyui_3/input",
# "output_dir": "waas/comfyui_3/output"
# },
# {
# "http_url": "https://bowongai-dev--comfyui-for-waas-ui-4.modal.run",
# "ws_url": "wss://bowongai-dev--comfyui-for-waas-ui-4.modal.run/ws",
# "input_dir": "waas/comfyui_4/input",
# "output_dir": "waas/comfyui_4/output"
# }
COMFYUI_SERVERS_JSON='[{"http_url":"http://127.0.0.1:8000","ws_url":"ws://127.0.0.1:8000/ws","input_dir":"waas/comfyui_4/input","output_dir":"waas/comfyui_4/output"}]'
DB_FILE="workflows_service.sqlite"

12
run_service.py
View File

@@ -4,6 +4,18 @@
"""
import sys
import os
from pathlib import Path
# 加载 .env 文件
try:
from dotenv import load_dotenv
env_path = Path(__file__).parent / ".env"
load_dotenv(env_path)
print(f"已加载环境变量文件: {env_path}")
except ImportError:
print("警告: 未安装 python-dotenv，无法加载 .env 文件")
except Exception as e:
print(f"警告: 加载 .env 文件失败: {e}")
# 将项目根目录添加到Python路径
project_root = os.path.dirname(os.path.abspath(__file__))

460
workflow_service/comfyui_client.py
View File

@@ -1,24 +1,211 @@
import asyncio
from collections import defaultdict
from ctypes import Union
import json
import logging
import os
import random
import uuid
from typing import Dict, Any
import websockets
from collections import defaultdict
from datetime import datetime
from typing import Dict, Any, Optional, List, Set
import aiohttp
import websockets
from aiohttp import ClientTimeout
from workflow_service.config import Settings, ComfyUIServer
from workflow_service.database import (
create_workflow_run,
update_workflow_run_status,
create_workflow_run_nodes,
update_workflow_run_node_status,
get_workflow_run,
get_pending_workflow_runs,
get_running_workflow_runs,
get_workflow_run_nodes,
)
settings = Settings()
API_INPUT_PREFIX = "INPUT_"
API_OUTPUT_PREFIX = "OUTPUT_"
# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
# [新增] 定义一个自定义异常，用于封装来自ComfyUI的执行错误
# 全局任务队列管理器
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
@@ -64,7 +251,7 @@ async def get_server_status(
except Exception as e:
# 当请求失败时，将返回上面定义的、结构正确的初始 status_info
print(f"警告: 无法检查服务器 {server.http_url} 的队列状态: {e}")
logger.warning(f"无法检查服务器 {server.http_url} 的队列状态: {e}")
return status_info
@@ -88,7 +275,7 @@ async def select_server_for_execution() -> ComfyUIServer:
if free_servers:
selected_server = random.choice(free_servers)
print(
logger.info(
f"发现 {len(free_servers)} 个空闲服务器。已选择: {selected_server.http_url}"
)
return selected_server
@@ -99,7 +286,7 @@ async def select_server_for_execution() -> ComfyUIServer:
]
if reachable_servers:
selected_server = random.choice(reachable_servers)
print(
logger.info(
f"所有服务器当前都在忙。从可达服务器中随机选择: {selected_server.http_url}"
)
return selected_server
@@ -108,17 +295,70 @@ async def select_server_for_execution() -> ComfyUIServer:
raise ConnectionError("所有配置的ComfyUI服务器都不可达。")
async def execute_prompt_on_server(prompt: Dict, server: ComfyUIServer) -> Dict:
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_id = await _queue_prompt(prompt, client_id, server.http_url)
print(f"工作流已在 {server.http_url} 上入队，Prompt ID: {prompt_id}")
results = await _get_execution_results(prompt_id, client_id, server.ws_url)
# 应用请求数据到工作流
patched_workflow = patch_workflow(workflow_data, api_spec, request_data)
# 转换为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参数名。
@@ -198,9 +438,13 @@ def parse_api_spec(workflow_data: dict) -> Dict[str, Dict[str, Any]]:
return spec
def patch_workflow(workflow_data: dict, api_spec: dict, request_data: dict) -> dict:
def patch_workflow(
workflow_data: dict,
api_spec: dict[str, dict[str, Any]],
request_data: dict[str, Any],
) -> dict:
"""
[BUG修复] 根据API请求数据，正确地修改工作流JSON中的 `widgets_values`。
将request_data中的参数值，patch到workflow_data中。并返回修改后的workflow_data。
"""
if "nodes" not in workflow_data:
raise ValueError("无效的工作流格式")
@@ -232,8 +476,8 @@ def patch_workflow(workflow_data: dict, api_spec: dict, request_data: dict) -> d
break
if target_widget_index == -1:
print(
f"警告: 在节点 {node_id} 中未找到名为 '{widget_name_to_patch}' 的 widget。跳过此参数。"
logger.warning(
f"在节点 {node_id} 中未找到名为 '{widget_name_to_patch}' 的 widget。跳过此参数。"
)
continue
@@ -259,8 +503,8 @@ def patch_workflow(workflow_data: dict, api_spec: dict, request_data: dict) -> d
try:
target_node["widgets_values"][target_widget_index] = target_type(value)
except (ValueError, TypeError) as e:
print(
f"警告: 无法将参数 '{param_name}' 的值 '{value}' 转换为类型 '{spec['type']}'。错误: {e}"
logger.warning(
f"无法将参数 '{param_name}' 的值 '{value}' 转换为类型 '{spec['type']}'。错误: {e}"
)
continue
@@ -322,25 +566,47 @@ def convert_workflow_to_prompt_api_format(workflow_data: dict) -> dict:
return prompt_api_format
async def _queue_prompt(prompt: dict, client_id: str, http_url: str) -> str:
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}
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"
async with session.post(prompt_url, json=payload) as response:
response.raise_for_status()
result = await response.json()
if "prompt_id" not in result:
raise Exception(f"从 ComfyUI /prompt 端点返回的响应无效: {result}")
return result["prompt_id"]
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(prompt_id: str, client_id: str, ws_url: str) -> dict:
async def _get_execution_results(
workflow_data: Dict,
prompt_id: str,
client_id: str,
ws_url: str,
workflow_run_id: str,
) -> dict:
"""
通过WebSocket连接到指定的ComfyUI服务器，聚合执行结果。
[核心改动] 新增对 'execution_error' 消息的处理。
现在支持节点级别的状态跟踪。
"""
full_ws_url = f"{ws_url}?clientId={client_id}"
aggregated_outputs = {}
@@ -360,34 +626,154 @@ async def _get_execution_results(prompt_id: str, client_id: str, ws_url: str) ->
if not (data and data.get("prompt_id") == prompt_id):
continue
# [核心改动] 捕获并处理执行错误
# 捕获并处理执行错误
if msg_type == "execution_error":
print(f"ComfyUI执行错误 (Prompt ID: {prompt_id}): {data}")
# 抛出自定义异常，将错误详情传递出去
raise ComfyUIExecutionError(data)
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 _get_execution_results_legacy(
prompt_id: str, client_id: str, ws_url: str
) -> dict:
"""
简化版的执行结果获取函数，不包含数据库状态跟踪。
"""
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}"
)
# 抛出自定义异常，将错误详情传递出去
raise ComfyUIExecutionError(error_data)
# 处理节点执行完成
if msg_type == "executed":
node_id = data.get("node")
output_data = data.get("output")
if node_id and output_data:
aggregated_outputs[node_id] = output_data
print(f"收到节点 {node_id} 的输出 (Prompt ID: {prompt_id})")
logger.info(
f"收到节点 {node_id} 的输出 (Prompt ID: {prompt_id})"
)
# 处理整个工作流执行完成
elif msg_type == "executing" and data.get("node") is None:
print(f"Prompt ID: {prompt_id} 执行完成。")
logger.info(f"Prompt ID: {prompt_id} 执行完成。")
return aggregated_outputs
except websockets.exceptions.ConnectionClosed as e:
print(f"WebSocket 连接已关闭 (Prompt ID: {prompt_id})。错误: {e}")
logger.warning(
f"WebSocket 连接已关闭 (Prompt ID: {prompt_id})。错误: {e}"
)
return aggregated_outputs
except Exception as e:
# 重新抛出我们自己的异常，或者处理其他意外错误
if not isinstance(e, ComfyUIExecutionError):
print(f"处理 prompt {prompt_id} 时发生意外错误: {e}")
logger.error(f"处理 prompt {prompt_id} 时发生意外错误: {e}")
raise e
except websockets.exceptions.InvalidURI as e:
print(
logger.error(
f"错误: 尝试连接的WebSocket URI无效: '{full_ws_url}'. 原始URL: '{ws_url}'. 错误: {e}"
)
raise e

180
workflow_service/database.py
View File

@@ -3,6 +3,7 @@ from typing import Optional
import aiosqlite
import json
import re
import uuid
from workflow_service.config import Settings
@@ -18,6 +19,43 @@ async def init_db():
workflow_json TEXT NOT NULL
)
""")
# 新增 workflow_run 表
await db.execute("""
CREATE TABLE IF NOT EXISTS workflow_run (
id TEXT PRIMARY KEY,
workflow_name TEXT NOT NULL,
prompt_id TEXT,
client_id TEXT,
status TEXT NOT NULL DEFAULT 'pending',
server_url TEXT,
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
started_at TIMESTAMP,
completed_at TIMESTAMP,
error_message TEXT,
workflow_json TEXT NOT NULL,
api_spec TEXT NOT NULL,
request_data TEXT NOT NULL,
result TEXT
)
""")
# 新增 workflow_run_nodes 表，记录每个节点的运行状态
await db.execute("""
CREATE TABLE IF NOT EXISTS workflow_run_nodes (
id TEXT PRIMARY KEY,
workflow_run_id TEXT NOT NULL,
node_id TEXT NOT NULL,
node_type TEXT NOT NULL,
status TEXT NOT NULL DEFAULT 'pending',
started_at TIMESTAMP,
completed_at TIMESTAMP,
output_data TEXT,
error_message TEXT,
FOREIGN KEY (workflow_run_id) REFERENCES workflow_run (id)
)
""")
await db.commit()
print(f"数据库 '{DATABASE_FILE}' 已准备就绪。")
@@ -72,4 +110,144 @@ async def delete_workflow(name: str) -> bool:
async with aiosqlite.connect(DATABASE_FILE) as db:
cursor = await db.execute("DELETE FROM workflows WHERE name = ?", (name,))
await db.commit()
return cursor.rowcount > 0
return cursor.rowcount > 0
# 新增 workflow_run 相关函数
async def create_workflow_run(
workflow_run_id: str,
workflow_name: str,
workflow_json: str,
api_spec: str,
request_data: str
) -> str:
"""创建新的工作流运行记录"""
async with aiosqlite.connect(DATABASE_FILE) as db:
await db.execute("""
INSERT INTO workflow_run (id, workflow_name, workflow_json, api_spec, request_data)
VALUES (?, ?, ?, ?, ?)
""", (workflow_run_id, workflow_name, workflow_json, api_spec, request_data))
await db.commit()
return workflow_run_id
async def update_workflow_run_status(
workflow_run_id: str,
status: str,
server_url: str = None,
prompt_id: str = None,
client_id: str = None,
error_message: str = None,
result: str = None
):
"""更新工作流运行状态"""
async with aiosqlite.connect(DATABASE_FILE) as db:
if status == 'running':
await db.execute("""
UPDATE workflow_run
SET status = ?, server_url = ?, prompt_id = ?, client_id = ?, started_at = CURRENT_TIMESTAMP
WHERE id = ?
""", (status, server_url, prompt_id, client_id, workflow_run_id))
elif status == 'completed':
await db.execute("""
UPDATE workflow_run
SET status = ?, completed_at = CURRENT_TIMESTAMP, result = ?
WHERE id = ?
""", (status, result, workflow_run_id))
elif status == 'failed':
await db.execute("""
UPDATE workflow_run
SET status = ?, error_message = ?, completed_at = CURRENT_TIMESTAMP
WHERE id = ?
""", (status, error_message, workflow_run_id))
else:
await db.execute("""
UPDATE workflow_run
SET status = ?
WHERE id = ?
""", (status, workflow_run_id))
await db.commit()
async def create_workflow_run_nodes(workflow_run_id: str, nodes_data: list):
"""创建工作流运行节点记录"""
async with aiosqlite.connect(DATABASE_FILE) as db:
for node in nodes_data:
node_id = str(uuid.uuid4())
await db.execute("""
INSERT INTO workflow_run_nodes (id, workflow_run_id, node_id, node_type)
VALUES (?, ?, ?, ?)
""", (node_id, workflow_run_id, node["id"], node["type"]))
await db.commit()
async def update_workflow_run_node_status(
workflow_run_id: str,
node_id: str,
status: str,
output_data: str = None,
error_message: str = None
):
"""更新工作流运行节点状态"""
async with aiosqlite.connect(DATABASE_FILE) as db:
if status == 'running':
await db.execute("""
UPDATE workflow_run_nodes
SET status = ?, started_at = CURRENT_TIMESTAMP
WHERE workflow_run_id = ? AND node_id = ?
""", (status, workflow_run_id, node_id))
elif status == 'completed':
await db.execute("""
UPDATE workflow_run_nodes
SET status = ?, output_data = ?, completed_at = CURRENT_TIMESTAMP
WHERE workflow_run_id = ? AND node_id = ?
""", (status, output_data, workflow_run_id, node_id))
elif status == 'failed':
await db.execute("""
UPDATE workflow_run_nodes
SET status = ?, error_message = ?, completed_at = CURRENT_TIMESTAMP
WHERE workflow_run_id = ? AND node_id = ?
""", (status, error_message, workflow_run_id, node_id))
else:
await db.execute("""
UPDATE workflow_run_nodes
SET status = ?
WHERE workflow_run_id = ? AND node_id = ?
""", (status, workflow_run_id, node_id))
await db.commit()
async def get_workflow_run(workflow_run_id: str) -> dict | None:
"""获取工作流运行记录"""
async with aiosqlite.connect(DATABASE_FILE) as db:
db.row_factory = aiosqlite.Row
cursor = await db.execute("""
SELECT * FROM workflow_run WHERE id = ?
""", (workflow_run_id,))
row = await cursor.fetchone()
return dict(row) if row else None
async def get_workflow_run_nodes(workflow_run_id: str) -> list[dict]:
"""获取工作流运行节点记录"""
async with aiosqlite.connect(DATABASE_FILE) as db:
db.row_factory = aiosqlite.Row
cursor = await db.execute("""
SELECT * FROM workflow_run_nodes WHERE workflow_run_id = ?
""", (workflow_run_id,))
rows = await cursor.fetchall()
return [dict(row) for row in rows]
async def get_pending_workflow_runs() -> list[dict]:
"""获取所有待处理的工作流运行记录"""
async with aiosqlite.connect(DATABASE_FILE) as db:
db.row_factory = aiosqlite.Row
cursor = await db.execute("""
SELECT * FROM workflow_run WHERE status = 'pending' ORDER BY created_at ASC
""")
rows = await cursor.fetchall()
return [dict(row) for row in rows]
async def get_running_workflow_runs() -> list[dict]:
"""获取所有正在运行的工作流运行记录"""
async with aiosqlite.connect(DATABASE_FILE) as db:
db.row_factory = aiosqlite.Row
cursor = await db.execute("""
SELECT * FROM workflow_run WHERE status = 'running'
""")
rows = await cursor.fetchall()
return [dict(row) for row in rows]

464
workflow_service/main.py
View File

@@ -15,7 +15,11 @@ from pydantic import BaseModel
from workflow_service import comfyui_client
from workflow_service import database
from workflow_service.utils.s3_client import upload_file_to_s3
from workflow_service.comfyui_client import ComfyUIExecutionError
from workflow_service.comfyui_client import (
ComfyUIExecutionError,
queue_manager,
submit_workflow_to_queue,
)
from workflow_service.config import Settings
settings = Settings()
@@ -81,7 +85,6 @@ async def startup_event():
print(f"错误: 无法在启动时初始化服务器目录: {e}")
# --- Section 1: 工作流管理API ---
@web_app.post("/api/workflow", status_code=201)
async def publish_workflow_endpoint(request: Request):
"""
@@ -117,7 +120,6 @@ async def get_all_workflows_endpoint():
raise HTTPException(status_code=500, detail=f"Failed to get workflows: {e}")
# [BUG修复] 修正API路径，使其与其他管理端点保持一致
@web_app.delete(f"/api/workflow/{{workflow_name:path}}")
async def delete_workflow_endpoint(
workflow_name: str = Path(
@@ -140,179 +142,8 @@ async def delete_workflow_endpoint(
raise HTTPException(status_code=500, detail=f"Failed to delete workflow: {e}")
# --- Section 2: 工作流执行API (核心改动) ---
def get_files_in_dir(directory: str) -> Set[str]:
file_set = set()
for root, _, files in os.walk(directory):
for filename in files:
if not filename.startswith("."):
file_set.add(os.path.join(root, filename))
return file_set
async def download_file_from_url(
session: aiohttp.ClientSession, url: str, save_path: str
):
async with session.get(url) as response:
response.raise_for_status()
with open(save_path, "wb") as f:
while True:
chunk = await response.content.read(8192)
if not chunk:
break
f.write(chunk)
async def handle_file_upload(file_path: str, base_name: str) -> str:
s3_object_name = f"outputs/{base_name}/{uuid.uuid4()}_{os.path.basename(file_path)}"
return await upload_file_to_s3(file_path, settings.S3_BUCKET_NAME, s3_object_name)
@web_app.post("/api/run/")
async def execute_workflow_endpoint(
base_name: str, request_data_raw: Dict[str, Any], version: Optional[str] = None
):
"""
执行工作流
"""
cleanup_paths = []
try:
# 1. 获取工作流定义
workflow_data = await database.get_workflow(base_name, version)
if not workflow_data:
detail = (
f"工作流 '{base_name}'"
+ (f" 带版本 '{version}'" if version else " (最新版)")
+ " 未找到。"
)
raise HTTPException(status_code=404, detail=detail)
workflow = json.loads(workflow_data["workflow_json"])
api_spec = comfyui_client.parse_api_spec(workflow)
request_data = {k.lower(): v for k, v in request_data_raw.items()}
# 2. [核心改动] 使用智能调度选择服务器
try:
selected_server = await comfyui_client.select_server_for_execution()
except Exception as e:
raise HTTPException(
status_code=503, detail=f"无法选择ComfyUI服务器执行任务: {e}"
)
server_input_dir = selected_server.input_dir
server_output_dir = selected_server.output_dir
# 3. 下载文件到选定服务器的输入目录
async with aiohttp.ClientSession() as session:
for param_name, param_def in api_spec["inputs"].items():
if param_def["type"] == "UploadFile" and param_name in request_data:
image_url = request_data[param_name]
if not isinstance(image_url, str) or not image_url.startswith(
"http"
):
raise HTTPException(
status_code=400,
detail=f"参数 '{param_name}' 必须是一个有效的URL。",
)
filename = f"api_download_{uuid.uuid4().hex}{os.path.splitext(image_url.split('/')[-1])[1] or '.dat'}"
save_path = os.path.join(server_input_dir, filename)
try:
await download_file_from_url(session, image_url, save_path)
# 直接更新 request_data，以便后续的 patch_workflow 使用
request_data[param_name] = filename
cleanup_paths.append(save_path)
except Exception as e:
raise HTTPException(
status_code=500,
detail=f"为 '{param_name}' 从 {image_url} 下载文件失败。错误: {e}",
)
# 4. Patch工作流，生成最终的API Prompt
patched_workflow = comfyui_client.patch_workflow(
workflow, api_spec, request_data
)
prompt_to_run = comfyui_client.convert_workflow_to_prompt_api_format(
patched_workflow
)
# 5. 执行前快照，并在选定服务器上执行工作流
files_before = get_files_in_dir(server_output_dir)
output_nodes = await comfyui_client.execute_prompt_on_server(
prompt_to_run, selected_server
)
# 6. 处理输出（与之前逻辑相同）
files_after = get_files_in_dir(server_output_dir)
new_files = files_after - files_before
output_response = {}
if new_files:
s3_urls = []
for file_path in new_files:
cleanup_paths.append(file_path)
try:
s3_urls.append(await handle_file_upload(file_path, base_name))
except Exception as e:
print(f"上传文件 {file_path} 到S3时出错: {e}")
if s3_urls:
output_response["output_files"] = s3_urls
for final_param_name, param_def in api_spec["outputs"].items():
node_id = param_def["node_id"]
if node_id in output_nodes:
node_output = output_nodes[node_id]
original_output_name = param_def["output_name"]
output_value = None
if original_output_name in node_output:
output_value = node_output[original_output_name]
elif "text" in node_output: # 备选，例如对于'ShowText'节点
output_value = node_output["text"]
if output_value is None:
continue
if isinstance(output_value, list):
output_value = output_value[0] if output_value else None
output_response[final_param_name] = output_value
return JSONResponse(content=output_response)
# [核心改动] 捕获来自ComfyUI的执行失败异常
except ComfyUIExecutionError as e:
print(f"捕获到ComfyUI执行错误: {e.error_data}")
# 返回 502 Bad Gateway 状态码，表示上游服务器出错
# detail 中包含结构化的错误信息，方便客户端处理
raise HTTPException(
status_code=500,
detail={
"message": "工作流在上游ComfyUI节点中执行失败。",
"error_details": e.error_data,
},
)
except Exception as e:
# 捕获其他所有异常，作为通用的服务器内部错误
print(f"执行工作流时发生未知错误: {e}")
raise HTTPException(status_code=500, detail=str(e))
finally:
if cleanup_paths:
print(f"正在清理 {len(cleanup_paths)} 个临时文件...")
for path in cleanup_paths:
try:
if os.path.exists(path):
os.remove(path)
print(f" - 已删除: {path}")
except Exception as e:
print(f" - 删除 {path} 时出错: {e}")
# --- Section 3: 工作流元数据/规范API ---
@web_app.get("/api/spec/")
async def get_workflow_spec_endpoint(base_name: str, version: Optional[str] = None):
@web_app.get("/api/workflow/{base_name:path}")
async def get_one_workflow_endpoint(base_name: str, version: Optional[str] = None):
"""
获取工作流规范
"""
@@ -328,14 +159,231 @@ async def get_workflow_spec_endpoint(base_name: str, version: Optional[str] = No
try:
workflow = json.loads(workflow_data["workflow_json"])
return comfyui_client.parse_api_spec(workflow)
return {
"workflow": workflow,
"api_spec": comfyui_client.parse_api_spec(workflow),
}
except Exception as e:
raise HTTPException(
status_code=500, detail=f"Failed to parse workflow specification: {e}"
)
# --- NEW Section: 服务器监控API ---
@web_app.post("/api/run")
async def run_workflow(
request: Request, workflow_name: str, workflow_version: Optional[str] = None
):
"""
异步执行工作流。
立即返回任务ID，调用者可以通过任务ID查询执行状态。
"""
try:
data = await request.json()
if not workflow_name:
raise HTTPException(status_code=400, detail="`workflow_name` 字段是必需的")
# 获取工作流定义
workflow_data = await database.get_workflow(workflow_name, workflow_version)
if not workflow_data:
detail = (
f"工作流 '{workflow_name}'"
+ (f" 带版本 '{workflow_version}'" if workflow_version else " (最新版)")
+ " 未找到。"
)
raise HTTPException(status_code=404, detail=detail)
workflow = json.loads(workflow_data["workflow_json"])
api_spec = comfyui_client.parse_api_spec(workflow)
# 提交到队列
workflow_run_id = await submit_workflow_to_queue(
workflow_name=workflow_name, workflow_data=workflow, api_spec=api_spec, request_data=data
)
return JSONResponse(
content={
"workflow_run_id": workflow_run_id,
"status": "queued",
"message": "工作流已提交到队列，正在等待执行",
},
status_code=202,
)
except Exception as e:
raise HTTPException(status_code=500, detail=f"提交工作流失败: {str(e)}")
@web_app.get("/api/run/{workflow_run_id}")
async def get_run_status(workflow_run_id: str):
"""
获取工作流执行状态。
"""
try:
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)}")
@web_app.get("/api/metrics")
async def get_metrics():
"""
获取队列状态概览。
"""
try:
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(queue_manager.running_tasks),
"queue_manager_status": "active",
}
except Exception as e:
raise HTTPException(status_code=500, detail=f"获取队列状态失败: {str(e)}")
# 同步工作流执行API（带S3上传功能）
# @web_app.post("/api/run_sync/")
# async def execute_workflow_sync_endpoint(
# base_name: str, request_data_raw: Dict[str, Any], version: Optional[str] = None
# ):
# """
# 同步执行工作流（支持S3文件上传）
# """
# cleanup_paths = []
# try:
# # 1. 获取工作流定义
# workflow_data = await database.get_workflow(base_name, version)
# if not workflow_data:
# detail = (
# f"工作流 '{base_name}'"
# + (f" 带版本 '{version}'" if version else " (最新版)")
# + " 未找到。"
# )
# raise HTTPException(status_code=404, detail=detail)
# workflow = json.loads(workflow_data["workflow_json"])
# api_spec = comfyui_client.parse_api_spec(workflow)
# request_data = {k.lower(): v for k, v in request_data_raw.items()}
# # 2. [核心改动] 使用智能调度选择服务器
# try:
# selected_server = await comfyui_client.select_server_for_execution()
# except Exception as e:
# raise HTTPException(
# status_code=503, detail=f"无法选择ComfyUI服务器执行任务: {e}"
# )
# server_input_dir = selected_server.input_dir
# server_output_dir = selected_server.output_dir
# # 3. 下载文件到选定服务器的输入目录
# async with aiohttp.ClientSession() as session:
# for param_name, param_def in api_spec["inputs"].items():
# if param_def["type"] == "UploadFile" and param_name in request_data:
# image_url = request_data[param_name]
# if not isinstance(image_url, str) or not image_url.startswith(
# "http"
# ):
# raise HTTPException(
# status_code=400,
# detail=f"参数 '{param_name}' 必须是一个有效的URL。",
# )
# filename = f"api_download_{uuid.uuid4().hex}{os.path.splitext(image_url.split('/')[-1])[1] or '.dat'}"
# save_path = os.path.join(server_input_dir, filename)
# try:
# await download_file_from_url(session, image_url, save_path)
# # 直接更新 request_data，以便后续的 patch_workflow 使用
# request_data[param_name] = filename
# cleanup_paths.append(save_path)
# except Exception as e:
# raise HTTPException(
# status_code=500,
# detail=f"为 '{param_name}' 从 {image_url} 下载文件失败。错误: {e}",
# )
# # 4. Patch工作流，生成最终的API Prompt
# patched_workflow = comfyui_client.patch_workflow(
# workflow, api_spec, request_data
# )
# prompt_to_run = comfyui_client.convert_workflow_to_prompt_api_format(
# patched_workflow
# )
# # 5. 执行前快照，并在选定服务器上执行工作流
# files_before = get_files_in_dir(server_output_dir)
# output_nodes = await comfyui_client.execute_prompt_on_server_legacy(
# prompt_to_run, selected_server
# )
# # 6. 处理输出（与之前逻辑相同）
# files_after = get_files_in_dir(server_output_dir)
# new_files = files_after - files_before
# output_response = {}
# if new_files:
# s3_urls = []
# for file_path in new_files:
# cleanup_paths.append(file_path)
# try:
# s3_urls.append(await handle_file_upload(file_path, base_name))
# except Exception as e:
# print(f"上传文件 {file_path} 到S3时出错: {e}")
# if s3_urls:
# output_response["output_files"] = s3_urls
# for final_param_name, param_def in api_spec["outputs"].items():
# node_id = param_def["node_id"]
# if node_id in output_nodes:
# node_output = output_nodes[node_id]
# original_output_name = param_def["output_name"]
# output_value = None
# if original_output_name in node_output:
# output_value = node_output[original_output_name]
# elif "text" in node_output: # 备选，例如对于'ShowText'节点
# output_value = node_output["text"]
# if output_value is None:
# continue
# if isinstance(output_value, list):
# output_value = output_value[0] if output_value else None
# output_response[final_param_name] = output_value
# return JSONResponse(content=output_response)
# # [核心改动] 捕获来自ComfyUI的执行失败异常
# except ComfyUIExecutionError as e:
# print(f"捕获到ComfyUI执行错误: {e.error_data}")
# # 返回 502 Bad Gateway 状态码，表示上游服务器出错
# # detail 中包含结构化的错误信息，方便客户端处理
# raise HTTPException(
# status_code=500,
# detail={
# "message": "工作流在上游ComfyUI节点中执行失败。",
# "error_details": e.error_data,
# },
# )
# except Exception as e:
# # 捕获其他所有异常，作为通用的服务器内部错误
# print(f"执行工作流时发生未知错误: {e}")
# raise HTTPException(status_code=500, detail=str(e))
# finally:
# if cleanup_paths:
# print(f"正在清理 {len(cleanup_paths)} 个临时文件...")
# for path in cleanup_paths:
# try:
# if os.path.exists(path):
# os.remove(path)
# print(f" - 已删除: {path}")
# except Exception as e:
# print(f" - 删除 {path} 时出错: {e}")
@web_app.get(
@@ -449,6 +497,27 @@ def read_root():
},
{
"step": 2,
"action": "异步执行工作流 (Execute Workflow Async with S3 Upload)",
"description": "提交工作流到异步队列，立即返回任务ID，支持长时间运行的工作流。执行完成后自动上传输出文件到S3。",
"endpoint": "/api/run",
"example_curl": 'curl -X POST "http://127.0.0.1:18000/api/run?workflow_name=my_t2i&workflow_version=20250801120000" -H \'Content-Type: application/json\' -d \'{\n "prompt_prompt": "a beautiful cat sitting on a roof",\n "sampler_seed": 12345\n}\'',
},
{
"step": 3,
"action": "查询异步任务状态 (Check Async Task Status)",
"description": "通过任务ID查询异步工作流的执行状态。",
"endpoint": "/api/run/{workflow_run_id}",
"example_curl": "curl -X GET 'http://127.0.0.1:18000/api/run/your-task-id-here'",
},
{
"step": 4,
"action": "查看队列状态 (Check Queue Status)",
"description": "查看当前异步任务队列的状态。",
"endpoint": "/api/metrics",
"example_curl": "curl -X GET 'http://127.0.0.1:18000/api/metrics'",
},
{
"step": 5,
"action": "查看工作流的API规范 (Inspect Workflow API Spec)",
"description": "一旦发布，你可以查询工作流的API规范，以了解需要提供哪些输入参数以及可以期望哪些输出。",
"endpoint": "/api/spec/",
@@ -465,10 +534,10 @@ def read_root():
"example_curl": "curl -X GET 'http://127.0.0.1:18000/api/spec/?base_name=my_t2i'",
},
{
"step": 3,
"action": "执行工作流 (Execute a Workflow)",
"description": "使用获取到的API规范，通过HTTP POST请求来执行工作流。对于文件输入，请提供可公开访问的URL。输出文件将被上传到S3并返回URL。",
"endpoint": "/api/run/",
"step": 6,
"action": "执行工作流 (Execute a Workflow - Sync with S3 Upload)",
"description": "使用获取到的API规范，通过HTTP POST请求来同步执行工作流。对于文件输入，请提供可公开访问的URL。输出文件将被自动上传到S3并返回URL。",
"endpoint": "/api/run_sync/",
"parameters": [
{"name": "base_name", "description": "工作流的基础名称。"},
{
@@ -476,10 +545,10 @@ def read_root():
"description": "可选。要执行的特定版本。如果省略，则执行最新版本。",
},
],
"example_curl": "curl -X POST 'http://127.0.0.1:18000/api/run/?base_name=my_t2i' -H 'Content-Type: application/json' -d '{\n \"prompt_prompt\": \"a beautiful cat sitting on a roof\",\n \"sampler_seed\": 12345\n}'",
"example_curl": "curl -X POST 'http://127.0.0.1:18000/api/run_sync/?base_name=my_t2i' -H 'Content-Type: application/json' -d '{\n \"prompt_prompt\": \"a beautiful cat sitting on a roof\",\n \"sampler_seed\": 12345\n}'",
},
{
"step": 4,
"step": 7,
"action": "管理工作流 (Manage Workflows)",
"description": "你也可以列出所有已发布的工作流或删除不再需要的工作流。",
"endpoints": [
@@ -503,6 +572,33 @@ def read_root():
return JSONResponse(content=guide)
def get_files_in_dir(directory: str) -> Set[str]:
file_set = set()
for root, _, files in os.walk(directory):
for filename in files:
if not filename.startswith("."):
file_set.add(os.path.join(root, filename))
return file_set
async def download_file_from_url(
session: aiohttp.ClientSession, url: str, save_path: str
):
async with session.get(url) as response:
response.raise_for_status()
with open(save_path, "wb") as f:
while True:
chunk = await response.content.read(8192)
if not chunk:
break
f.write(chunk)
async def handle_file_upload(file_path: str, base_name: str) -> str:
s3_object_name = f"outputs/{base_name}/{uuid.uuid4()}_{os.path.basename(file_path)}"
return await upload_file_to_s3(file_path, settings.S3_BUCKET_NAME, s3_object_name)
if __name__ == "__main__":
uvicorn.run(
"workflow_service.main:web_app", host="0.0.0.0", port=18000, reload=True