diff --git a/workflow_service/comfy/comfy_queue.py b/workflow_service/comfy/comfy_queue.py new file mode 100644 index 0000000..71ee2c8 --- /dev/null +++ b/workflow_service/comfy/comfy_queue.py @@ -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 diff --git a/workflow_service/comfy/comfy_workflow.py b/workflow_service/comfy/comfy_workflow.py new file mode 100644 index 0000000..31810c7 --- /dev/null +++ b/workflow_service/comfy/comfy_workflow.py @@ -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" diff --git a/workflow_service/comfyui_client.py b/workflow_service/comfyui_client.py deleted file mode 100644 index 14ae034..0000000 --- a/workflow_service/comfyui_client.py +++ /dev/null @@ -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() diff --git a/workflow_service/config.py b/workflow_service/config.py index f83f977..e01b425 100644 --- a/workflow_service/config.py +++ b/workflow_service/config.py @@ -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() \ No newline at end of file +settings = Settings() diff --git a/workflow_service/routes/run.py b/workflow_service/routes/run.py index 53c7b0f..4b48c0d 100644 --- a/workflow_service/routes/run.py +++ b/workflow_service/routes/run.py @@ -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)}") diff --git a/workflow_service/routes/runx/model_with_multi_dress.py b/workflow_service/routes/runx/model_with_multi_dress.py index e5395e8..5571e1f 100644 --- a/workflow_service/routes/runx/model_with_multi_dress.py +++ b/workflow_service/routes/runx/model_with_multi_dress.py @@ -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, diff --git a/workflow_service/routes/service.py b/workflow_service/routes/service.py index ff8c6b9..4284058 100644 --- a/workflow_service/routes/service.py +++ b/workflow_service/routes/service.py @@ -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) diff --git a/workflow_service/routes/workflow.py b/workflow_service/routes/workflow.py index 20c7e2c..190410b 100644 --- a/workflow_service/routes/workflow.py +++ b/workflow_service/routes/workflow.py @@ -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(