diff --git a/workflow_service/comfyui_client.py b/workflow_service/comfyui_client.py index 01a5566..5f0ffac 100644 --- a/workflow_service/comfyui_client.py +++ b/workflow_service/comfyui_client.py @@ -1,4 +1,6 @@ import asyncio +from collections import defaultdict +from ctypes import Union import json import random import uuid @@ -12,6 +14,9 @@ from workflow_service.config import Settings, ComfyUIServer settings = Settings() +API_INPUT_PREFIX = "INPUT_" +API_OUTPUT_PREFIX = "OUTPUT_" + # [新增] 定义一个自定义异常，用于封装来自ComfyUI的执行错误 class ComfyUIExecutionError(Exception): @@ -114,6 +119,209 @@ async def execute_prompt_on_server(prompt: Dict, server: ComfyUIServer) -> Dict: return results +def parse_api_spec(workflow_data: dict) -> Dict[str, Dict[str, Any]]: + """ + 解析工作流，并根据规范 '{基础名}_{属性名}_{可选计数}' 生成API参数名。 + """ + 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"]} + input_name_counter = defaultdict(int) + output_name_counter = defaultdict(int) + + for node_id, node in nodes_map.items(): + title: str = node.get("title") + if not title: + continue + + 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", []): + 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() + 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" + + spec["inputs"][final_param_name] = { + "node_id": node_id, + "type": param_type, + "widget_name": a_input["widget"]["name"], + } + + 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 + ) + + 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), + } + + return spec + + +def patch_workflow(workflow_data: dict, api_spec: dict, request_data: dict) -> dict: + """ + [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 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: + """ + 将工作流（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(prompt: dict, client_id: str, http_url: str) -> str: """通过HTTP POST将工作流任务提交到指定的ComfyUI服务器。""" for node_id in prompt: diff --git a/workflow_service/database.py b/workflow_service/database.py index 020c720..93c2c57 100644 --- a/workflow_service/database.py +++ b/workflow_service/database.py @@ -1,4 +1,5 @@ from datetime import datetime +from typing import Optional import aiosqlite import json import re @@ -54,7 +55,6 @@ async def get_latest_workflow_by_base_name(base_name: str) -> dict | None: return dict(row) if row else None async def get_workflow_by_version(base_name: str, version: str) -> dict | None: - # [BUG修复] 修复了工作流名称拼接错误 name = f"{base_name} [{version}]" async with aiosqlite.connect(DATABASE_FILE) as db: db.row_factory = aiosqlite.Row @@ -62,6 +62,12 @@ async def get_workflow_by_version(base_name: str, version: str) -> dict | None: row = await cursor.fetchone() return dict(row) if row else None +async def get_workflow(name: str, version: Optional[str] = None) -> dict | None: + if version: + return await get_workflow_by_version(name, version) + else: + return await get_latest_workflow_by_base_name(name) + 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,)) diff --git a/workflow_service/main.py b/workflow_service/main.py index 8cc7da1..9baa004 100644 --- a/workflow_service/main.py +++ b/workflow_service/main.py @@ -14,8 +14,7 @@ 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.utils.s3_client import upload_file_to_s3 from workflow_service.comfyui_client import ComfyUIExecutionError from workflow_service.config import Settings @@ -35,6 +34,7 @@ web_app.add_middleware( # --- Pydantic Response Models for New Endpoints --- + class ServerQueueDetails(BaseModel): running_count: int pending_count: int @@ -91,11 +91,15 @@ async def publish_workflow_endpoint(request: Request): 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.") + 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=201) + 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: @@ -115,8 +119,12 @@ async def get_all_workflows_endpoint(): # [BUG修复] 修正API路径，使其与其他管理端点保持一致 @web_app.delete(f"/api/workflow/{{workflow_name:path}}") -async def delete_workflow_endpoint(workflow_name: str = Path(..., - description="The full, unique name of the workflow to delete, e.g., 'my_workflow [20250101120000]'")): +async def delete_workflow_endpoint( + workflow_name: str = Path( + ..., + description="The full, unique name of the workflow to delete, e.g., 'my_workflow [20250101120000]'", + ) +): """ 删除工作流 """ @@ -125,74 +133,90 @@ async def delete_workflow_endpoint(workflow_name: str = Path(..., if success: return {"status": "deleted", "name": workflow_name} else: - raise HTTPException(status_code=404, detail=f"Workflow '{workflow_name}' 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('.'): + 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 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: + with open(save_path, "wb") as f: while True: chunk = await response.content.read(8192) - if not chunk: break + 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 s3_client.upload_file_to_s3(file_path, settings.S3_BUCKET_NAME, s3_object_name) + 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): +async def execute_workflow_endpoint( + base_name: str, request_data_raw: Dict[str, Any], version: Optional[str] = None +): """ 执行工作流 """ cleanup_paths = [] try: # 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) + workflow_data = await database.get_workflow(base_name, version) if not workflow_data: - detail = f"工作流 '{base_name}'" + (f" 带版本 '{version}'" if version else " (最新版)") + " 未找到。" + 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) + 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}") + 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: + 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。") + 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) @@ -202,16 +226,24 @@ async def execute_workflow_endpoint(base_name: str, request_data_raw: Dict[str, 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}") + raise HTTPException( + status_code=500, + detail=f"为 '{param_name}' 从 {image_url} 下载文件失败。错误: {e}", + ) # 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) + 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) + output_nodes = await comfyui_client.execute_prompt_on_server( + prompt_to_run, selected_server + ) # 6. 处理输出（与之前逻辑相同） files_after = get_files_in_dir(server_output_dir) @@ -229,11 +261,11 @@ async def execute_workflow_endpoint(base_name: str, request_data_raw: Dict[str, if s3_urls: output_response["output_files"] = s3_urls - for final_param_name, spec in api_spec["outputs"].items(): - node_id = spec["node_id"] + 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 = spec["output_name"] + original_output_name = param_def["output_name"] output_value = None if original_output_name in node_output: @@ -241,7 +273,8 @@ async def execute_workflow_endpoint(base_name: str, request_data_raw: Dict[str, elif "text" in node_output: # 备选，例如对于'ShowText'节点 output_value = node_output["text"] - if output_value is None: continue + if output_value is None: + continue if isinstance(output_value, list): output_value = output_value[0] if output_value else None @@ -258,8 +291,8 @@ async def execute_workflow_endpoint(base_name: str, request_data_raw: Dict[str, status_code=500, detail={ "message": "工作流在上游ComfyUI节点中执行失败。", - "error_details": e.error_data - } + "error_details": e.error_data, + }, ) except Exception as e: # 捕获其他所有异常，作为通用的服务器内部错误 @@ -283,25 +316,31 @@ async def get_workflow_spec_endpoint(base_name: str, version: Optional[str] = No """ 获取工作流规范 """ - 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) + workflow_data = await database.get_workflow(base_name, version) if not workflow_data: - detail = f"Workflow '{base_name}'" + (f" with version '{version}'" if version else " (latest)") + " 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) + workflow = json.loads(workflow_data["workflow_json"]) + return comfyui_client.parse_api_spec(workflow) except Exception as e: - raise HTTPException(status_code=500, detail=f"Failed to parse workflow specification: {e}") + raise HTTPException( + status_code=500, detail=f"Failed to parse workflow specification: {e}" + ) # --- NEW Section: 服务器监控API --- -@web_app.get("/api/servers/status", response_model=List[ServerStatus], tags=["Server Monitoring"]) + +@web_app.get( + "/api/servers/status", response_model=List[ServerStatus], tags=["Server Monitoring"] +) async def get_servers_status_endpoint(): """ 获取所有已配置的ComfyUI服务器的配置信息和实时状态。 @@ -311,7 +350,9 @@ async def get_servers_status_endpoint(): return [] async with aiohttp.ClientSession() as session: - status_tasks = [comfyui_client.get_server_status(server, session) for server in servers] + status_tasks = [ + comfyui_client.get_server_status(server, session) for server in servers + ] live_statuses = await asyncio.gather(*status_tasks) response_list = [] @@ -326,7 +367,7 @@ async def get_servers_status_endpoint(): output_dir=server_config.output_dir, is_reachable=status_data["is_reachable"], is_free=status_data["is_free"], - queue_details=status_data["queue_details"] + queue_details=status_data["queue_details"], ) ) return response_list @@ -347,7 +388,7 @@ async def _get_folder_contents(path: str) -> List[FileDetails]: FileDetails( name=entry.name, size_kb=round(stat.st_size / 1024, 2), - modified_at=datetime.fromtimestamp(stat.st_mtime) + modified_at=datetime.fromtimestamp(stat.st_mtime), ) ) except OSError as e: @@ -357,28 +398,36 @@ async def _get_folder_contents(path: str) -> List[FileDetails]: 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="服务器在配置列表中的索引")): +@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}。") + 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) + _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 + output_files=output_files, ) @@ -396,7 +445,7 @@ def read_root(): "action": "发布一个工作流 (Publish a Workflow)", "description": "从 ComfyUI 保存你的工作流 (API格式)，然后使用一个唯一的名称将其发布到服务中。名称中必须包含一个版本号，格式为 `[YYYYMMDDHHMMSS]`。", "endpoint": f"/api/workflow", - "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}'" + "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, @@ -404,11 +453,16 @@ def read_root(): "description": "一旦发布，你可以查询工作流的API规范，以了解需要提供哪些输入参数以及可以期望哪些输出。", "endpoint": "/api/spec/", "parameters": [ - {"name": "base_name", "description": "工作流的基础名称 (不含版本部分)。"}, - {"name": "version", - "description": "可选。工作流的具体版本号 (YYYYMMDDHHMMSS)。如果省略，则获取最新版本。"} + { + "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'" + "example_curl": "curl -X GET 'http://127.0.0.1:18000/api/spec/?base_name=my_t2i'", }, { "step": 3, @@ -417,27 +471,39 @@ def read_root(): "endpoint": "/api/run/", "parameters": [ {"name": "base_name", "description": "工作流的基础名称。"}, - {"name": "version", "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}'" + "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"/api/workflow", "description": "列出所有工作流。"}, - {"method": "DELETE", "path": f"/api/workflow/{{workflow_name}}", - "description": "按完整名称删除一个工作流。"} + { + "method": "GET", + "path": f"/api/workflow", + "description": "列出所有工作流。", + }, + { + "method": "DELETE", + "path": f"/api/workflow/{{workflow_name}}", + "description": "按完整名称删除一个工作流。", + }, ], "example_curl_list": f"curl -X GET http://127.0.0.1:18000/api/workflow", - "example_curl_delete": f"curl -X DELETE 'http://127.0.0.1:18000/api/workflow/my_t2i%20[20250801120000]'" - } + "example_curl_delete": f"curl -X DELETE 'http://127.0.0.1:18000/api/workflow/my_t2i%20[20250801120000]'", + }, ], - "notes": "请确保在 curl 命令中对包含空格或特殊字符的工作流名称进行URL编码。" + "notes": "请确保在 curl 命令中对包含空格或特殊字符的工作流名称进行URL编码。", } return JSONResponse(content=guide) if __name__ == "__main__": - uvicorn.run("workflow_service.main:web_app", host="0.0.0.0", port=18000, reload=True) + uvicorn.run( + "workflow_service.main:web_app", host="0.0.0.0", port=18000, reload=True + ) diff --git a/workflow_service/s3_client.py b/workflow_service/utils/s3_client.py similarity index 100% rename from workflow_service/s3_client.py rename to workflow_service/utils/s3_client.py diff --git a/workflow_service/workflow_parser.py b/workflow_service/workflow_parser.py deleted file mode 100644 index 475530f..0000000 --- a/workflow_service/workflow_parser.py +++ /dev/null @@ -1,185 +0,0 @@ -import re -from collections import defaultdict - -API_INPUT_PREFIX = "INPUT_" -API_OUTPUT_PREFIX = "OUTPUT_" - - -def parse_api_spec(workflow_data: dict) -> dict: - """ - 解析工作流，并根据规范 '{基础名}_{属性名}_{可选计数}' 生成API参数名。 - """ - 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"]} - input_name_counter = defaultdict(int) - output_name_counter = defaultdict(int) - - for node_id, node in nodes_map.items(): - title = node.get("title") - if not title: continue - - 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", []): - 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() - 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" - - spec["inputs"][final_param_name] = { - "node_id": node_id, - "type": param_type, - "widget_name": a_input["widget"]["name"] - } - - 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 - - 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) - } - - return spec - - -def patch_workflow(workflow_data: dict, api_spec: dict, request_data: dict) -> dict: - """ - [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 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: - """ - 将工作流（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 \ No newline at end of file