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

fix

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

464
workflow_service/main.py
View File

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