From db5d97cc495ddacb20c3317b36ef513116483bf8 Mon Sep 17 00:00:00 2001 From: "kyj@bowong.ai" Date: Fri, 1 Aug 2025 14:32:57 +0800 Subject: [PATCH] =?UTF-8?q?add=20waas=E6=9C=8D=E5=8A=A1=E5=99=A8=E6=94=AF?= =?UTF-8?q?=E6=8C=81=E5=A4=9Acomfyui=E6=89=A7=E8=A1=8C=E5=99=A8?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .env | 35 ++- README.md | 2 + modal_deploy.py | 8 +- workflow_service/comfyui_client.py | 141 ++++++++--- workflow_service/config.py | 47 +++- workflow_service/database.py | 90 +++---- workflow_service/main.py | 375 ++++++++++++++++++++-------- workflow_service/workflow_parser.py | 166 +++++++----- 8 files changed, 605 insertions(+), 259 deletions(-) diff --git a/.env b/.env index 9f82efb..801663a 100644 --- a/.env +++ b/.env @@ -1,10 +1,31 @@ -# ComfyUI服务器地址 -COMFYUI_URL="ws://bowongai--local-ui-4.modal.run/ws" -COMFYUI_HTTP_URL="https://bowongai--local-ui-4.modal.run" -# 绝对路径！例如: /home/user/ComfyUI/input -COMFYUI_INPUT_DIR="/db/input" -# 绝对路径！例如: /home/user/ComfyUI/output -COMFYUI_OUTPUT_DIR="/db/output" +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" + } +]' + +DB_FILE="/db/workflows_service.sqlite" # AWS S3 配置 S3_BUCKET_NAME="modal-media-cache" diff --git a/README.md b/README.md index 08ac422..6d59f13 100644 --- a/README.md +++ b/README.md @@ -10,6 +10,8 @@ # WAAS(工作流即服务) Demo API服务器 - 路径: ./workflow_service +- 部署: ```modal deploy modal_deploy.py --name waas-demo -e dev``` +- 配置: 参考```.env```文件 - 必须按照指定规则命名 - **输入节点名**: 前缀 **INPUT_** - **除生成文件节点外输出节点名**: 前缀 **OUTPUT_** diff --git a/modal_deploy.py b/modal_deploy.py index 5e1bc9d..b40a88b 100644 --- a/modal_deploy.py +++ b/modal_deploy.py @@ -12,6 +12,7 @@ fastapi_image = ( app = modal.App(image=fastapi_image, name="waas-demo") vol = modal.Volume.from_name("comfy_model", environment_name="dev", create_if_missing=True) +secret = modal.Secret.from_name("aws-s3-secret", environment_name="dev") with fastapi_image.imports(): @@ -20,7 +21,12 @@ with fastapi_image.imports(): memory=(128, 4096), scaledown_window=1200, volumes={ - "/db": vol + "/db": vol, + "/waas": modal.CloudBucketMount( + bucket_name="modal-media-cache", + secret=secret, + key_prefix="waas/" + ), } ) @modal.concurrent(max_inputs=100) diff --git a/workflow_service/comfyui_client.py b/workflow_service/comfyui_client.py index bd28a56..8d926d4 100644 --- a/workflow_service/comfyui_client.py +++ b/workflow_service/comfyui_client.py @@ -1,38 +1,110 @@ -import websockets +import asyncio import json -import uuid -import aiohttp import random -from workflow_service.config import Settings +import uuid +from typing import Dict, Any + +import aiohttp +import websockets +from aiohttp import ClientTimeout + +from workflow_service.config import Settings, ComfyUIServer + settings = Settings() -async def queue_prompt(prompt: dict, client_id: str) -> str: - """通过HTTP POST将工作流任务提交到ComfyUI队列。""" - # [缓存破解] 在每个节点中添加一个随机数输入，强制重新执行 +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 + print(f"警告: 无法检查服务器 {server.http_url} 的队列状态: {e}") + + return status_info + + +async def select_server_for_execution() -> ComfyUIServer: + """ + 智能选择一个ComfyUI服务器。 + 优先选择一个空闲的服务器，如果所有服务器都忙，则随机选择一个。 + """ + servers = settings.SERVERS + if not servers: + raise ValueError("没有在 COMFYUI_SERVERS_JSON 中配置任何服务器。") + if len(servers) == 1: + return servers[0] + + async with aiohttp.ClientSession() as session: + tasks = [get_server_status(server, session) for server in servers] + results = await asyncio.gather(*tasks) + + free_servers = [servers[i] for i, status in enumerate(results) if status["is_free"]] + + if free_servers: + selected_server = random.choice(free_servers) + print(f"发现 {len(free_servers)} 个空闲服务器。已选择: {selected_server.http_url}") + return selected_server + else: + # 后备方案：选择一个可达的服务器，即使它很忙 + reachable_servers = [servers[i] for i, status in enumerate(results) if status["is_reachable"]] + if reachable_servers: + selected_server = random.choice(reachable_servers) + print(f"所有服务器当前都在忙。从可达服务器中随机选择: {selected_server.http_url}") + return selected_server + else: + # 最坏情况：所有服务器都不可达，抛出异常 + raise ConnectionError("所有配置的ComfyUI服务器都不可达。") + + +async def queue_prompt(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} - async with aiohttp.ClientSession() as session: - http_url = f"{settings.COMFYUI_HTTP_URL}/prompt" - async with session.post(http_url, json=payload) as response: + 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"Invalid response from ComfyUI /prompt endpoint: {result}") + raise Exception(f"从 ComfyUI /prompt 端点返回的响应无效: {result}") return result["prompt_id"] -async def get_execution_results(prompt_id: str, client_id: str) -> dict: - """ - 通过WebSocket连接，聚合所有'executed'事件的输出， - 直到整个执行流程结束。 - """ - ws_url = f"{settings.COMFYUI_URL}?clientId={client_id}" +async def get_execution_results(prompt_id: str, client_id: str, ws_url: str) -> dict: + """通过WebSocket连接到指定的ComfyUI服务器，聚合执行结果。""" + full_ws_url = f"{ws_url}?clientId={client_id}" aggregated_outputs = {} - - async with websockets.connect(ws_url) as websocket: + async with websockets.connect(full_ws_url) as websocket: while True: try: out = await websocket.recv() @@ -40,36 +112,31 @@ async def get_execution_results(prompt_id: str, client_id: str) -> dict: message = json.loads(out) msg_type = message.get('type') data = message.get('data') - - # 我们只关心与我们prompt_id相关的事件 if data and data.get('prompt_id') == prompt_id: 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"Output received for node {node_id} (Prompt ID: {prompt_id})") - - # 判断执行是否结束 - # 官方UI使用 "executing" 且 node is null 作为结束标志 + print(f"收到节点 {node_id} 的输出 (Prompt ID: {prompt_id})") elif msg_type == 'executing' and data.get('node') is None: - print(f"Execution finished for Prompt ID: {prompt_id}") + print(f"Prompt ID: {prompt_id} 执行完成。") return aggregated_outputs - except websockets.exceptions.ConnectionClosed as e: - print(f"WebSocket connection closed for {prompt_id}. Returning aggregated results. Error: {e}") - return aggregated_outputs # 连接关闭也视为结束 + print(f"WebSocket 连接已关闭 (Prompt ID: {prompt_id})。错误: {e}") + return aggregated_outputs except Exception as e: - print(f"An error occurred for {prompt_id}: {e}") + print(f"处理 prompt {prompt_id} 时发生错误: {e}") break return aggregated_outputs -async def run_workflow(prompt: dict) -> dict: - """主协调函数：提交任务，然后等待结果。""" +async def execute_prompt_on_server(prompt: Dict, server: ComfyUIServer) -> Dict: + """ + 在指定的服务器上执行一个准备好的prompt。 + """ client_id = str(uuid.uuid4()) - prompt_id = await queue_prompt(prompt, client_id) - print(f"Workflow successfully queued with Prompt ID: {prompt_id}") - results = await get_execution_results(prompt_id, client_id) + 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) return results \ No newline at end of file diff --git a/workflow_service/config.py b/workflow_service/config.py index f632b55..f83f977 100644 --- a/workflow_service/config.py +++ b/workflow_service/config.py @@ -1,15 +1,48 @@ -from dotenv import dotenv_values +import json +from pydantic import BaseModel, ValidationError from pydantic_settings import BaseSettings, SettingsConfigDict +from typing import List + + +class ComfyUIServer(BaseModel): + """定义单个ComfyUI服务器的配置模型。""" + http_url: str + ws_url: str + input_dir: str + output_dir: str + class Settings(BaseSettings): - model_config = SettingsConfigDict(env_file='./.env', env_file_encoding='utf-8') - COMFYUI_URL: str = "ws://127.0.0.1:8188/ws" - COMFYUI_HTTP_URL: str = "http://127.0.0.1:8188" - COMFYUI_INPUT_DIR: str - COMFYUI_OUTPUT_DIR: str + """ + 应用配置类，通过.env文件加载环境变量。 + """ + 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"}]' + + DB_FILE: str = "/db/workflows_service.sqlite" + # S3 和 AWS 配置保持不变 S3_BUCKET_NAME: str AWS_ACCESS_KEY_ID: str AWS_SECRET_ACCESS_KEY: str AWS_REGION_NAME: str -# settings = Settings() \ No newline at end of file + @property + def SERVERS(self) -> List[ComfyUIServer]: + """ + 解析JSON配置字符串，返回一个经过验证的服务器配置对象列表。 + """ + try: + server_list = json.loads(self.COMFYUI_SERVERS_JSON) + # 使用Pydantic模型进行验证和类型转换 + return [ComfyUIServer(**server) for server in server_list] + except json.JSONDecodeError: + raise ValueError("COMFYUI_SERVERS_JSON 格式无效，必须是有效的JSON数组。") + except ValidationError as e: + raise ValueError(f"COMFYUI_SERVERS_JSON 中的服务器配置无效: {e}") + except Exception as e: + raise e + + +settings = Settings() \ No newline at end of file diff --git a/workflow_service/database.py b/workflow_service/database.py index 8641c51..3611b81 100644 --- a/workflow_service/database.py +++ b/workflow_service/database.py @@ -2,58 +2,62 @@ import aiosqlite import json import re -DATABASE_FILE = "/db/workflows_service.sqlite" +from workflow_service.config import Settings + +DATABASE_FILE = Settings().DB_FILE async def init_db(): - async with aiosqlite.connect(DATABASE_FILE) as db: - await db.execute(""" - CREATE TABLE IF NOT EXISTS workflows ( - name TEXT PRIMARY KEY, - base_name TEXT NOT NULL, - version TEXT NOT NULL, - workflow_json TEXT NOT NULL - ) - """) - await db.commit() - print(f"数据库 '{DATABASE_FILE}' 已准备就绪。") + async with aiosqlite.connect(DATABASE_FILE) as db: + await db.execute(""" + CREATE TABLE IF NOT EXISTS workflows ( + name TEXT PRIMARY KEY, + base_name TEXT NOT NULL, + version TEXT NOT NULL, + workflow_json TEXT NOT NULL + ) + """) + await db.commit() + print(f"数据库 '{DATABASE_FILE}' 已准备就绪。") async def save_workflow(name: str, workflow_json: str): - version_match = re.search(r"\[(20\d{12})\]$", name) - if not version_match: - raise ValueError("Workflow name must have a version suffix like [YYYYMMDDHHMMSS]") - version = version_match.group(1) - base_name = re.sub(r"\s*\[(20\d{12})\]$", "", name).strip() - async with aiosqlite.connect(DATABASE_FILE) as db: - await db.execute( - "INSERT OR REPLACE INTO workflows (name, base_name, version, workflow_json) VALUES (?, ?, ?, ?)", - (name, base_name, version, workflow_json) - ) - await db.commit() + # 正则表达式修正为匹配更精确的版本格式 [YYYYMMDDHHMMSS] + version_match = re.search(r"\[(\d{14})\]$", name) + if not version_match: + raise ValueError("Workflow name must have a version suffix like [YYYYMMDDHHMMSS]") + version = version_match.group(1) + base_name = re.sub(r"\s*\[\d{14}\]$", "", name).strip() + async with aiosqlite.connect(DATABASE_FILE) as db: + await db.execute( + "INSERT OR REPLACE INTO workflows (name, base_name, version, workflow_json) VALUES (?, ?, ?, ?)", + (name, base_name, version, workflow_json) + ) + await db.commit() async def get_all_workflows() -> list[dict]: - async with aiosqlite.connect(DATABASE_FILE) as db: - db.row_factory = aiosqlite.Row - cursor = await db.execute("SELECT name, workflow_json FROM workflows") - rows = await cursor.fetchall() - return [{"name": row["name"], "workflow": json.loads(row["workflow_json"])} for row in rows] + async with aiosqlite.connect(DATABASE_FILE) as db: + db.row_factory = aiosqlite.Row + cursor = await db.execute("SELECT name, workflow_json FROM workflows") + rows = await cursor.fetchall() + return [{"name": row["name"], "workflow": json.loads(row["workflow_json"])} for row in rows] async def get_latest_workflow_by_base_name(base_name: str) -> dict | None: - async with aiosqlite.connect(DATABASE_FILE) as db: - db.row_factory = aiosqlite.Row - cursor = await db.execute("SELECT * FROM workflows WHERE base_name = ? ORDER BY version DESC LIMIT 1", (base_name,)) - row = await cursor.fetchone() - return dict(row) if row else None + async with aiosqlite.connect(DATABASE_FILE) as db: + db.row_factory = aiosqlite.Row + cursor = await db.execute("SELECT * FROM workflows WHERE base_name = ? ORDER BY version DESC LIMIT 1", (base_name,)) + row = await cursor.fetchone() + return dict(row) if row else None async def get_workflow_by_version(base_name: str, version: str) -> dict | None: - name = f"{base_name} [{version}]" - async with aiosqlite.connect(DATABASE_FILE) as db: - db.row_factory = aiosqlite.Row - cursor = await db.execute("SELECT * FROM workflows WHERE name = ?", (name,)) - row = await cursor.fetchone() - return dict(row) if row else None + # [BUG修复] 修复了工作流名称拼接错误 + name = f"{base_name} [{version}]" + async with aiosqlite.connect(DATABASE_FILE) as db: + db.row_factory = aiosqlite.Row + cursor = await db.execute("SELECT * FROM workflows WHERE name = ?", (name,)) + row = await cursor.fetchone() + return dict(row) if row else None 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 \ No newline at end of file + 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 \ No newline at end of file diff --git a/workflow_service/main.py b/workflow_service/main.py index 81d8db4..0b3fdeb 100644 --- a/workflow_service/main.py +++ b/workflow_service/main.py @@ -1,43 +1,92 @@ -import aiohttp -from workflow_service import comfyui_client -from workflow_service import database +import asyncio import json import os -from workflow_service import s3_client import uuid -from workflow_service import workflow_parser +from datetime import datetime from typing import Optional, List, Dict, Any, Set +import aiohttp import uvicorn from fastapi import FastAPI, Request, HTTPException, Path from fastapi.responses import JSONResponse +from pydantic import BaseModel +from workflow_service import comfyui_client +from workflow_service import database +from workflow_service import s3_client +from workflow_service import workflow_parser from workflow_service.config import Settings + settings = Settings() web_app = FastAPI(title="ComfyUI Workflow Service & Management API") +# --- Pydantic Response Models for New Endpoints --- + +class ServerQueueDetails(BaseModel): + running_count: int + pending_count: int + + +class ServerStatus(BaseModel): + server_index: int + http_url: str + ws_url: str + input_dir: str + output_dir: str + is_reachable: bool + is_free: bool + queue_details: ServerQueueDetails + + +class FileDetails(BaseModel): + name: str + size_kb: float + modified_at: datetime + + +class ServerFiles(BaseModel): + server_index: int + http_url: str + input_files: List[FileDetails] + output_files: List[FileDetails] + + @web_app.on_event("startup") -async def startup_event(): await database.init_db(); os.makedirs(settings.COMFYUI_INPUT_DIR, - exist_ok=True); os.makedirs( - settings.COMFYUI_OUTPUT_DIR, exist_ok=True) +async def startup_event(): + """服务启动时，初始化数据库并为所有配置的服务器创建输入/输出目录。""" + await database.init_db() + try: + servers = settings.SERVERS + print(f"检测到 {len(servers)} 个 ComfyUI 服务器配置。") + for server in servers: + print(f" - 正在为服务器 {server.http_url} 准备目录...") + os.makedirs(server.input_dir, exist_ok=True) + print(f" - 输入目录: {os.path.abspath(server.input_dir)}") + os.makedirs(server.output_dir, exist_ok=True) + print(f" - 输出目录: {os.path.abspath(server.output_dir)}") + except ValueError as e: + print(f"错误: 无法在启动时初始化服务器目录: {e}") -# --- Section 1: 工作流管理API (无改动) --- -# ... (代码与上一版完全相同) ... +# --- Section 1: 工作流管理API --- BASE_MANAGEMENT_PATH = "/api/workflow" -@web_app.post(BASE_MANAGEMENT_PATH, status_code=200) +@web_app.post(BASE_MANAGEMENT_PATH, status_code=201) async def publish_workflow_endpoint(request: Request): try: - data = await request.json(); - name, wf_json = data.get("name"), data.get( - "workflow"); - await database.save_workflow(name, json.dumps(wf_json)); + data = await request.json() + name, wf_json = data.get("name"), data.get("workflow") + if not name or not wf_json: + raise HTTPException(status_code=400, detail="`name` and `workflow` fields are required.") + await database.save_workflow(name, json.dumps(wf_json)) + print(f"Workflow '{name}' published.") return JSONResponse( - content={"status": "success", "message": f"Workflow '{name}' published."}, status_code=200) + content={"status": "success", "message": f"Workflow '{name}' published."}, status_code=201) + except ValueError as e: + raise HTTPException(status_code=400, detail=str(e)) except Exception as e: raise HTTPException(status_code=500, detail=f"Failed to save workflow: {e}") @@ -50,167 +99,161 @@ async def get_all_workflows_endpoint(): raise HTTPException(status_code=500, detail=f"Failed to get workflows: {e}") +# [BUG修复] 修正API路径，使其与其他管理端点保持一致 @web_app.delete(f"{BASE_MANAGEMENT_PATH}/{{workflow_name:path}}") -async def delete_workflow_endpoint(workflow_name: str = Path(..., title="...")): +async def delete_workflow_endpoint(workflow_name: str = Path(..., + description="The full, unique name of the workflow to delete, e.g., 'my_workflow [20250101120000]'")): try: - success = await database.delete_workflow(workflow_name); + success = await database.delete_workflow(workflow_name) if success: - return {"status": "deleted", "name": workflow_name}; + return {"status": "deleted", "name": workflow_name} else: - raise HTTPException(status_code=404, detail="Workflow not found") + raise HTTPException(status_code=404, detail=f"Workflow '{workflow_name}' not found") except Exception as e: 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)) + 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(); + response.raise_for_status() with open(save_path, 'wb') as f: while True: - chunk = await response.content.read(8192); - if not chunk: - break; + 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并返回URL""" s3_object_name = f"outputs/{base_name}/{uuid.uuid4()}_{os.path.basename(file_path)}" return await s3_client.upload_file_to_s3(file_path, settings.S3_BUCKET_NAME, s3_object_name) -@web_app.post("/api/run/{base_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. 获取工作流和处理输入 (与上一版相同) - # ... + # 1. 获取工作流定义 if version: workflow_data = await database.get_workflow_by_version(base_name, version) else: workflow_data = await database.get_latest_workflow_by_base_name(base_name) - if not workflow_data: raise HTTPException(status_code=404, detail=f"Workflow '{base_name}' not found.") + + 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 = workflow_parser.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, spec in api_spec["inputs"].items(): if spec["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"Parameter '{param_name}' must be a valid URL.") - original_filename = image_url.split('/')[-1].split('?')[0]; - _, file_extension = os.path.splitext(original_filename) - if not file_extension: file_extension = '.dat' - filename = f"api_download_{uuid.uuid4()}{file_extension}" - save_path = os.path.join(settings.COMFYUI_INPUT_DIR, filename) + 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[param_name] = filename; + 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"Failed to download file for '{param_name}' from {image_url}. Error: {e}") + detail=f"为 '{param_name}' 从 {image_url} 下载文件失败。错误: {e}") - # 2. 执行前快照 - files_before = get_files_in_dir(settings.COMFYUI_OUTPUT_DIR) - - # 3. Patch, 转换并执行 (缓存破解已移入client) + # 4. Patch工作流，生成最终的API Prompt patched_workflow = workflow_parser.patch_workflow(workflow, api_spec, request_data) prompt_to_run = workflow_parser.convert_workflow_to_prompt_api_format(patched_workflow) - output_nodes = await comfyui_client.run_workflow(prompt_to_run) - # 4. 执行后快照并计算差异 - files_after = get_files_in_dir(settings.COMFYUI_OUTPUT_DIR) + # 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 - # 5. [核心修正] 统一处理所有输出 output_response = {} - processed_files = set() # 记录已通过文件系统快照处理的文件 - - # 5.1 首先处理所有新生成的文件 if new_files: s3_urls = [] for file_path in new_files: cleanup_paths.append(file_path) - processed_files.add(os.path.basename(file_path)) try: s3_urls.append(await handle_file_upload(file_path, base_name)) except Exception as e: - print(f"Error uploading file {file_path} to S3: {e}") - if s3_urls: output_response["output_files"] = s3_urls + print(f"上传文件 {file_path} 到S3时出错: {e}") + if s3_urls: + output_response["output_files"] = s3_urls - # 5.2 然后处理WebSocket返回的非文件输出，并检查文本输出是否是文件路径 for final_param_name, spec in api_spec["outputs"].items(): node_id = spec["node_id"] if node_id in output_nodes: node_output = output_nodes[node_id] original_output_name = spec["output_name"] + output_value = None if original_output_name in node_output: output_value = node_output[original_output_name] - # 展开列表 - if isinstance(output_value, list): - output_value = output_value[0] if output_value else None - - # 检查文本输出是否是未被发现的文件路径 - if isinstance(output_value, str) and ( - '.png' in output_value or '.jpg' in output_value or '.mp4' in output_value or 'output' in output_value): - potential_filename = os.path.basename(output_value.replace('\\', '/')) - if potential_filename not in processed_files: - # 这是一个新的文件路径，尝试在output目录中找到它 - potential_path = os.path.join(settings.COMFYUI_OUTPUT_DIR, potential_filename) - if os.path.exists(potential_path): - print(f"Found extra file from text output: {potential_path}") - cleanup_paths.append(potential_path) - processed_files.add(potential_filename) - try: - s3_url = await handle_file_upload(potential_path, base_name) - # 将它也加入output_files列表 - if "output_files" not in output_response: output_response["output_files"] = [] - output_response["output_files"].append(s3_url) - except Exception as e: - print(f"Error uploading extra file {potential_path} to S3: {e}") - continue # 处理完毕，跳过将其作为文本输出 - output_response[final_param_name] = output_value - elif "text" in node_output: + elif "text" in node_output: # 备选，例如对于'ShowText'节点 output_value = node_output["text"] - # 如果不是文件，则作为普通值输出 - output_response[final_param_name] = output_value - return output_response + 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) finally: - # 清理操作保持不变 - print(f"Cleaning up {len(cleanup_paths)} temporary files...") - for path in cleanup_paths: - try: - if os.path.exists(path): os.remove(path); print(f" - Deleted: {path}") - except Exception as e: - print(f" - Error deleting {path}: {e}") + 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/{base_name}") +# --- Section 3: 工作流元数据/规范API --- +@web_app.get("/api/spec/") async def get_workflow_spec_endpoint(base_name: str, version: Optional[str] = None): - # ... if version: workflow_data = await database.get_workflow_by_version(base_name, version) else: workflow_data = await database.get_latest_workflow_by_base_name(base_name) + if not workflow_data: - detail = f"Workflow '{base_name}'" + (f" with version '{version}'" if version else "") + " not found." + detail = f"Workflow '{base_name}'" + (f" with version '{version}'" if version else " (latest)") + " not found." raise HTTPException(status_code=404, detail=detail) + try: workflow = json.loads(workflow_data['workflow_json']) return workflow_parser.parse_api_spec(workflow) @@ -218,9 +261,145 @@ async def get_workflow_spec_endpoint(base_name: str, version: Optional[str] = No raise HTTPException(status_code=500, detail=f"Failed to parse workflow specification: {e}") -@web_app.get("/") -def read_root(): return {"message": "Welcome to the ComfyUI Workflow Service API!"} +# --- NEW Section: 服务器监控API --- + +@web_app.get("/api/servers/status", response_model=List[ServerStatus], tags=["Server Monitoring"]) +async def get_servers_status_endpoint(): + """ + 获取所有已配置的ComfyUI服务器的配置信息和实时状态。 + """ + servers = settings.SERVERS + if not servers: + return [] + + async with aiohttp.ClientSession() as session: + status_tasks = [comfyui_client.get_server_status(server, session) for server in servers] + live_statuses = await asyncio.gather(*status_tasks) + + response_list = [] + for i, server_config in enumerate(servers): + status_data = live_statuses[i] + response_list.append( + ServerStatus( + server_index=i, + http_url=server_config.http_url, + ws_url=server_config.ws_url, + input_dir=server_config.input_dir, + output_dir=server_config.output_dir, + is_reachable=status_data["is_reachable"], + is_free=status_data["is_free"], + queue_details=status_data["queue_details"] + ) + ) + return response_list + + +async def _get_folder_contents(path: str) -> List[FileDetails]: + """异步地列出并返回文件夹内容的详细信息。""" + if not os.path.isdir(path): + return [] + + def sync_list_files(dir_path): + files = [] + try: + for entry in os.scandir(dir_path): + if entry.is_file(): + stat = entry.stat() + files.append( + FileDetails( + name=entry.name, + size_kb=round(stat.st_size / 1024, 2), + modified_at=datetime.fromtimestamp(stat.st_mtime) + ) + ) + except OSError as e: + print(f"无法扫描目录 {dir_path}: {e}") + return sorted(files, key=lambda x: x.modified_at, reverse=True) + + return await asyncio.to_thread(sync_list_files, path) + + +@web_app.get("/api/servers/{server_index}/files", response_model=ServerFiles, tags=["Server Monitoring"]) +async def list_server_files_endpoint(server_index: int = Path(..., ge=0, description="服务器在配置列表中的索引")): + """ + 获取指定ComfyUI服务器的输入和输出文件夹中的文件列表。 + """ + servers = settings.SERVERS + if server_index >= len(servers): + raise HTTPException(status_code=404, + detail=f"服务器索引 {server_index} 超出范围。有效索引为 0 到 {len(servers) - 1}。") + + server_config = servers[server_index] + + input_files, output_files = await asyncio.gather( + _get_folder_contents(server_config.input_dir), + _get_folder_contents(server_config.output_dir) + ) + + return ServerFiles( + server_index=server_index, + http_url=server_config.http_url, + input_files=input_files, + output_files=output_files + ) + + +@web_app.get("/", response_class=JSONResponse) +def read_root(): + """ + 提供一个API的快速使用指南。 + """ + guide = { + "service_name": "ComfyUI Workflow Service & Management API", + "description": "一个用于发布、管理和执行 ComfyUI 工作流的API服务。它将 ComfyUI 的图形化工作流抽象为标准的 RESTful API 端点。", + "quick_start_guide": [ + { + "step": 1, + "action": "发布一个工作流 (Publish a Workflow)", + "description": "从 ComfyUI 保存你的工作流 (API格式)，然后使用一个唯一的名称将其发布到服务中。名称中必须包含一个版本号，格式为 `[YYYYMMDDHHMMSS]`。", + "endpoint": f"POST {BASE_MANAGEMENT_PATH}", + "example_curl": "curl -X POST http://127.0.0.1:18000/api/workflow -H 'Content-Type: application/json' -d '{\n \"name\": \"my_t2i [20250801120000]\",\n \"workflow\": { ... your_workflow_api_json ... }\n}'" + }, + { + "step": 2, + "action": "查看工作流的API规范 (Inspect Workflow API Spec)", + "description": "一旦发布，你可以查询工作流的API规范，以了解需要提供哪些输入参数以及可以期望哪些输出。", + "endpoint": "/api/spec/", + "parameters": [ + {"name": "base_name", "description": "工作流的基础名称 (不含版本部分)。"}, + {"name": "version", + "description": "可选。工作流的具体版本号 (YYYYMMDDHHMMSS)。如果省略，则获取最新版本。"} + ], + "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/", + "parameters": [ + {"name": "base_name", "description": "工作流的基础名称。"}, + {"name": "version", "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}'" + }, + { + "step": 4, + "action": "管理工作流 (Manage Workflows)", + "description": "你也可以列出所有已发布的工作流或删除不再需要的工作流。", + "endpoints": [ + {"method": "GET", "path": f"{BASE_MANAGEMENT_PATH}", "description": "列出所有工作流。"}, + {"method": "DELETE", "path": f"{BASE_MANAGEMENT_PATH}/{{workflow_name}}", + "description": "按完整名称删除一个工作流。"} + ], + "example_curl_list": f"curl -X GET http://127.0.0.1:18000{BASE_MANAGEMENT_PATH}", + "example_curl_delete": f"curl -X DELETE 'http://127.0.0.1:18000/api/workflow/my_t2i%20[20250801120000]'" + } + ], + "notes": "请确保在 curl 命令中对包含空格或特殊字符的工作流名称进行URL编码。" + } + return JSONResponse(content=guide) if __name__ == "__main__": - uvicorn.run(web_app, host="127.0.0.1", port=18000) + uvicorn.run("workflow_service.main:web_app", host="0.0.0.0", port=18000, reload=True) diff --git a/workflow_service/workflow_parser.py b/workflow_service/workflow_parser.py index f4e8c22..475530f 100644 --- a/workflow_service/workflow_parser.py +++ b/workflow_service/workflow_parser.py @@ -14,8 +14,6 @@ def parse_api_spec(workflow_data: dict) -> dict: 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_name_counter = defaultdict(int) output_name_counter = defaultdict(int) @@ -23,56 +21,44 @@ def parse_api_spec(workflow_data: dict) -> dict: title = node.get("title") if not title: continue - # --- 处理API输入 --- if title.startswith(API_INPUT_PREFIX): base_name = title[len(API_INPUT_PREFIX):].lower() if "inputs" in node: for a_input in node.get("inputs", []): - # 只关心由用户直接控制的widget输入 if a_input.get("link") is None and "widget" in a_input: widget_name = a_input["widget"]["name"].lower() - # 构建基础参数名 + # [BUG修复] 构建有意义的基础参数名 param_name_candidate = f"{base_name}_{widget_name}" - # 处理重名 input_name_counter[param_name_candidate] += 1 count = input_name_counter[param_name_candidate] final_param_name = f"{param_name_candidate}_{count}" if count > 1 else param_name_candidate - # 类型推断 input_type_str = a_input.get("type", "STRING").upper() - if "COMBO" in a_input.get("type"): - # 如果是加载类节点，名字为'image'或'video'的widget是文件上传 - if a_input["widget"]["name"] in ["image", "video"]: - param_type = "UploadFile" - else: - param_type = "string" # 加载节点其他参数默认为string + param_type = "string" + if "COMBO" in input_type_str: + param_type = "UploadFile" elif "INT" in input_type_str: param_type = "int" elif "FLOAT" in input_type_str: param_type = "float" - else: - param_type = "string" spec["inputs"][final_param_name] = { "node_id": node_id, "type": param_type, - "widget_name": a_input["widget"]["name"] # 保留原始widget名用于patch + "widget_name": a_input["widget"]["name"] } - # --- 处理API输出 --- elif title.startswith(API_OUTPUT_PREFIX): base_name = title[len(API_OUTPUT_PREFIX):].lower() - if "outputs" in node: for an_output in node.get("outputs", []): output_name = an_output["name"].lower() - # 构建基础参数名 + # [BUG修复] 构建有意义的基础参数名 param_name_candidate = f"{base_name}_{output_name}" - # 处理重名 output_name_counter[param_name_candidate] += 1 count = output_name_counter[param_name_candidate] final_param_name = f"{param_name_candidate}_{count}" if count > 1 else param_name_candidate @@ -80,72 +66,120 @@ def parse_api_spec(workflow_data: dict) -> dict: spec["outputs"][final_param_name] = { "node_id": node_id, "class_type": node.get("type"), - "output_name": an_output["name"], # 保留原始输出名用于结果匹配 - "output_index": node["outputs"].index(an_output) # 保留索引 + "output_name": an_output["name"], + "output_index": node["outputs"].index(an_output) } return spec -# --- 其他函数保持不变 --- - def patch_workflow(workflow_data: dict, api_spec: dict, request_data: dict) -> dict: - # (此函数与上一版完全相同，无需改动) - if "nodes" not in workflow_data: raise ValueError("Invalid workflow format") + """ + [BUG修复] 根据API请求数据，正确地修改工作流JSON中的 `widgets_values`。 + """ + if "nodes" not in workflow_data: + raise ValueError("无效的工作流格式") + nodes_map = {str(node["id"]): node for node in workflow_data["nodes"]} + for param_name, value in request_data.items(): - if param_name in api_spec["inputs"]: - spec = api_spec["inputs"][param_name] - node_id = spec["node_id"] - if node_id not in nodes_map: continue - target_node = nodes_map[node_id] - widget_name_to_patch = spec["widget_name"] - widgets_values = target_node.get("widgets_values", {}) - if isinstance(widgets_values, dict): - widgets_values[widget_name_to_patch] = value - elif isinstance(widgets_values, list): - widget_cursor = 0 - for input_config in target_node.get("inputs", []): - if "widget" in input_config: - if input_config["widget"].get("name") == widget_name_to_patch: - if widget_cursor < len(widgets_values): - target_type = str - if spec['type'] == 'int': - target_type = int - elif spec['type'] == 'float': - target_type = float - widgets_values[widget_cursor] = target_type(value) - break - widget_cursor += 1 - target_node["widgets_values"] = widgets_values + if param_name not in api_spec["inputs"]: + continue + + spec = api_spec["inputs"][param_name] + node_id = spec["node_id"] + if node_id not in nodes_map: + continue + + target_node = nodes_map[node_id] + widget_name_to_patch = 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: + print(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 spec['type'] == 'int': + target_type = int + elif spec['type'] == 'float': + target_type = float + + # 在正确的位置上更新值 + try: + target_node["widgets_values"][target_widget_index] = target_type(value) + except (ValueError, TypeError) as e: + print(f"警告: 无法将参数 '{param_name}' 的值 '{value}' 转换为类型 '{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: - # (此函数与上一版完全相同，无需改动) - if "nodes" not in workflow_data: raise ValueError("Invalid workflow format") - prompt_api_format, link_map = {}, {} - for link in workflow_data.get("links", []): - link_map[link[0]] = [str(link[1]), link[2]] + """ + 将工作流（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", []) - if isinstance(widgets_values, dict): - for key, val in widgets_values.items(): - if not isinstance(val, dict): inputs_dict[key] = val - elif isinstance(widgets_values, list): - widget_idx_counter = 0 - for input_config in node.get("inputs", []): - if "widget" in input_config: - if widget_idx_counter < len(widgets_values): - inputs_dict[input_config["name"]] = widgets_values[widget_idx_counter] - widget_idx_counter += 1 + 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} + + prompt_api_format[node_id] = { + "class_type": node["type"], + "inputs": inputs_dict + } + return prompt_api_format \ No newline at end of file