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

add .gitignore, update README with new API server details, and add requirements.txt for dependencies

Browse Source
This commit is contained in:
iHeyTang
2025-08-08 17:39:20 +08:00
parent 6cc9aeab32
commit ab12c1482d
4 changed files with 79 additions and 43 deletions
Download Patch File Download Diff File Expand all files Collapse all files

72
workflow_service/comfyui_client.py
View File

@@ -12,6 +12,7 @@ from workflow_service.config import Settings, ComfyUIServer
settings = Settings()
# [新增] 定义一个自定义异常，用于封装来自ComfyUI的执行错误
class ComfyUIExecutionError(Exception):
def __init__(self, error_data: dict):
@@ -24,7 +25,10 @@ class ComfyUIExecutionError(Exception):
)
super().__init__(message)
async def get_server_status(server: ComfyUIServer, session: aiohttp.ClientSession) -> Dict[str, Any]:
async def get_server_status(
server: ComfyUIServer, session: aiohttp.ClientSession
) -> Dict[str, Any]:
"""
检查单个ComfyUI服务器的详细状态。
返回一个包含可达性、队列状态和详细队列内容的字典。
@@ -33,10 +37,7 @@ async def get_server_status(server: ComfyUIServer, session: aiohttp.ClientSessio
status_info = {
"is_reachable": False,
"is_free": False,
"queue_details": {
"running_count": 0,
"pending_count": 0
}
"queue_details": {"running_count": 0, "pending_count": 0},
}
try:
queue_url = f"{server.http_url}/queue"
@@ -51,10 +52,10 @@ async def get_server_status(server: ComfyUIServer, session: aiohttp.ClientSessio
status_info["queue_details"] = {
"running_count": running_count,
"pending_count": pending_count
"pending_count": pending_count,
}
status_info["is_free"] = (running_count == 0 and pending_count == 0)
status_info["is_free"] = running_count == 0 and pending_count == 0
except Exception as e:
# 当请求失败时，将返回上面定义的、结构正确的初始 status_info
@@ -82,21 +83,38 @@ async def select_server_for_execution() -> ComfyUIServer:
if free_servers:
selected_server = random.choice(free_servers)
print(f"发现 {len(free_servers)} 个空闲服务器。已选择: {selected_server.http_url}")
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"]]
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}")
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:
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, 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
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()
@@ -111,7 +129,7 @@ async def queue_prompt(prompt: dict, client_id: str, http_url: str) -> str:
return result["prompt_id"]
async def get_execution_results(prompt_id: str, client_id: str, ws_url: str) -> dict:
async def _get_execution_results(prompt_id: str, client_id: str, ws_url: str) -> dict:
"""
通过WebSocket连接到指定的ComfyUI服务器，聚合执行结果。
[核心改动] 新增对 'execution_error' 消息的处理。
@@ -128,26 +146,26 @@ async def get_execution_results(prompt_id: str, client_id: str, ws_url: str) ->
continue
message = json.loads(out)
msg_type = message.get('type')
data = message.get('data')
msg_type = message.get("type")
data = message.get("data")
if not (data and data.get('prompt_id') == prompt_id):
if not (data and data.get("prompt_id") == prompt_id):
continue
# [核心改动] 捕获并处理执行错误
if msg_type == 'execution_error':
if msg_type == "execution_error":
print(f"ComfyUI执行错误 (Prompt ID: {prompt_id}): {data}")
# 抛出自定义异常，将错误详情传递出去
raise ComfyUIExecutionError(data)
if msg_type == 'executed':
node_id = data.get('node')
output_data = data.get('output')
if msg_type == "executed":
node_id = data.get("node")
output_data = data.get("output")
if node_id and output_data:
aggregated_outputs[node_id] = output_data
print(f"收到节点 {node_id} 的输出 (Prompt ID: {prompt_id})")
elif msg_type == 'executing' and data.get('node') is None:
elif msg_type == "executing" and data.get("node") is None:
print(f"Prompt ID: {prompt_id} 执行完成。")
return aggregated_outputs
@@ -161,17 +179,9 @@ async def get_execution_results(prompt_id: str, client_id: str, ws_url: str) ->
raise e
except websockets.exceptions.InvalidURI as e:
print(f"错误: 尝试连接的WebSocket URI无效: '{full_ws_url}'. 原始URL: '{ws_url}'. 错误: {e}")
print(
f"错误: 尝试连接的WebSocket URI无效: '{full_ws_url}'. 原始URL: '{ws_url}'. 错误: {e}"
)
raise e
return aggregated_outputs
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, 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

7
workflow_service/requirements.txt
View File

@@ -1,7 +0,0 @@
fastapi
uvicorn[standard]
pydantic-settings
websockets
aiohttp
boto3
aiosqlite