refactor: workflow_service改名app
This commit is contained in:
iHeyTang
234
app/comfy/comfy_queue.py
Normal file
234
app/comfy/comfy_queue.py
Normal file
@@ -0,0 +1,234 @@
|
||||
import asyncio
|
||||
import logging
|
||||
from typing import Any, Optional
|
||||
|
||||
import aiohttp
|
||||
|
||||
from app.comfy.comfy_workflow import ComfyWorkflow
|
||||
from app.comfy.comfy_run import ComfyRun
|
||||
from app.config import Settings
|
||||
from app.comfy.comfy_server import server_manager, ComfyUIServerInfo
|
||||
|
||||
settings = Settings()
|
||||
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
# 全局任务队列管理器
|
||||
class WorkflowQueueManager:
|
||||
def __init__(self, monitor_interval: int = 5):
|
||||
self.running_tasks = {} # server_url -> task_info
|
||||
self.pending_tasks = [] # 等待队列
|
||||
self.lock = asyncio.Lock()
|
||||
self._queue_monitor_task: Optional[asyncio.Task] = None
|
||||
self._monitor_interval = monitor_interval # 监控间隔（秒）
|
||||
self._last_processing_time = None # 上次处理队列的时间
|
||||
|
||||
async def start_queue_monitor(self):
|
||||
"""启动队列监控任务"""
|
||||
if self._queue_monitor_task is None or self._queue_monitor_task.done():
|
||||
self._queue_monitor_task = asyncio.create_task(self._monitor_queue())
|
||||
logger.info(f"队列监控任务已启动，监控间隔: {self._monitor_interval}秒")
|
||||
|
||||
async def stop_queue_monitor(self):
|
||||
"""停止队列监控任务"""
|
||||
if self._queue_monitor_task and not self._queue_monitor_task.done():
|
||||
self._queue_monitor_task.cancel()
|
||||
try:
|
||||
await self._queue_monitor_task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
logger.info("队列监控任务已停止")
|
||||
|
||||
async def add_task(
|
||||
self,
|
||||
workflow: ComfyWorkflow,
|
||||
request_data: dict,
|
||||
):
|
||||
"""添加新任务到队列"""
|
||||
async with self.lock:
|
||||
# 创建ComfyRun实例并保存到数据库
|
||||
comfy_run = await ComfyRun.create(workflow, request_data)
|
||||
|
||||
# 添加到待处理队列
|
||||
self.pending_tasks.append(comfy_run.run_id)
|
||||
logger.info(
|
||||
f"任务 {comfy_run.run_id} 已添加到队列，当前队列长度: {len(self.pending_tasks)}"
|
||||
)
|
||||
|
||||
return comfy_run.run_id
|
||||
|
||||
async def get_server_status(
|
||||
self, server: ComfyUIServerInfo, 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 _monitor_queue(self):
|
||||
"""定期监控队列，当有可用服务器时自动处理"""
|
||||
while True:
|
||||
try:
|
||||
await asyncio.sleep(self._monitor_interval)
|
||||
|
||||
# 每次定时检测触发时，打印当前状态
|
||||
async with self.lock:
|
||||
from datetime import datetime
|
||||
|
||||
pending_count = len(self.pending_tasks)
|
||||
running_count = len(self.running_tasks)
|
||||
current_time = datetime.now()
|
||||
timestamp = current_time.strftime("%Y-%m-%d %H:%M:%S")
|
||||
|
||||
# 计算距离上次处理队列的时间
|
||||
time_since_last_processing = ""
|
||||
if self._last_processing_time:
|
||||
time_diff = current_time - self._last_processing_time
|
||||
minutes = int(time_diff.total_seconds() // 60)
|
||||
seconds = int(time_diff.total_seconds() % 60)
|
||||
time_since_last_processing = (
|
||||
f"上次处理: {minutes}分{seconds}秒前"
|
||||
)
|
||||
else:
|
||||
time_since_last_processing = "上次处理: 从未"
|
||||
|
||||
# 格式化状态信息，提高可读性
|
||||
status_info = (
|
||||
f"⏰ 定时检测触发 [{timestamp}]\t"
|
||||
f"待处理/运行中: {pending_count}/{running_count}\t"
|
||||
f"监控间隔: {self._monitor_interval} 秒\t"
|
||||
f"🕐 {time_since_last_processing}"
|
||||
)
|
||||
logger.info(status_info)
|
||||
|
||||
if self.pending_tasks:
|
||||
available_servers = await self._get_available_servers()
|
||||
if available_servers:
|
||||
logger.info(
|
||||
f"监控发现 {len(available_servers)} 个可用服务器，开始处理队列"
|
||||
)
|
||||
# 使用create_task避免阻塞监控循环
|
||||
asyncio.create_task(self._process_queue())
|
||||
else:
|
||||
# 只在有任务等待时记录日志，避免日志噪音
|
||||
if len(self.pending_tasks) > 0:
|
||||
logger.debug(
|
||||
f"队列中有 {len(self.pending_tasks)} 个待处理任务，但无可用服务器"
|
||||
)
|
||||
else:
|
||||
# 队列为空时，减少日志输出
|
||||
logger.debug("队列为空，继续监控中...")
|
||||
|
||||
except asyncio.CancelledError:
|
||||
logger.info("队列监控任务被取消")
|
||||
break
|
||||
except Exception as e:
|
||||
logger.error(f"队列监控任务出错: {e}")
|
||||
# 出错时稍微延长等待时间，避免频繁重试
|
||||
await asyncio.sleep(self._monitor_interval * 2)
|
||||
|
||||
async def _process_queue(self):
|
||||
"""处理队列中的任务"""
|
||||
async with self.lock:
|
||||
# 更新上次处理队列的时间
|
||||
from datetime import datetime
|
||||
|
||||
self._last_processing_time = datetime.now()
|
||||
|
||||
if not self.pending_tasks:
|
||||
logger.debug("队列为空，无需处理")
|
||||
return
|
||||
|
||||
# 检查是否有空闲的服务器
|
||||
available_servers = await self._get_available_servers()
|
||||
if not available_servers:
|
||||
logger.info(
|
||||
f"队列中有 {len(self.pending_tasks)} 个待处理任务，但没有可用的服务器"
|
||||
)
|
||||
return
|
||||
|
||||
# 获取一个待处理任务
|
||||
workflow_run_id = self.pending_tasks.pop(0)
|
||||
server = available_servers[0]
|
||||
|
||||
logger.info(
|
||||
f"开始处理任务 {workflow_run_id}，使用服务器 {server.name} ({server.http_url})"
|
||||
)
|
||||
|
||||
# 记录运行中任务
|
||||
self.running_tasks[server.http_url] = {
|
||||
"workflow_run_id": workflow_run_id,
|
||||
"started_at": datetime.now(),
|
||||
"server_name": server.name,
|
||||
}
|
||||
|
||||
# 启动任务执行
|
||||
asyncio.create_task(self._execute_task(workflow_run_id, server))
|
||||
|
||||
logger.info(
|
||||
f"任务 {workflow_run_id} 已分配到服务器 {server.name}，当前运行中任务数: {len(self.running_tasks)}"
|
||||
)
|
||||
|
||||
async def _get_available_servers(self) -> list[ComfyUIServerInfo]:
|
||||
"""获取可用的服务器"""
|
||||
# 使用新的服务器管理器获取可用服务器
|
||||
available_server = await server_manager.get_available_server()
|
||||
if available_server:
|
||||
return [available_server]
|
||||
return []
|
||||
|
||||
async def _execute_task(self, workflow_run_id: str, server: ComfyUIServerInfo):
|
||||
"""执行任务"""
|
||||
try:
|
||||
# 从run_id创建ComfyRun实例并执行
|
||||
comfy_run = await ComfyRun.from_run_id(workflow_run_id)
|
||||
if not comfy_run:
|
||||
raise Exception(f"找不到工作流运行记录: {workflow_run_id}")
|
||||
|
||||
# 执行工作流（ComfyRun内部处理所有逻辑包括资源管理）
|
||||
await comfy_run.execute(server)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"执行任务 {workflow_run_id} 时出错: {e}")
|
||||
|
||||
finally:
|
||||
# 清理运行状态
|
||||
async with self.lock:
|
||||
if server.http_url in self.running_tasks:
|
||||
del self.running_tasks[server.http_url]
|
||||
|
||||
|
||||
# 全局队列管理器实例
|
||||
# 从配置中读取监控间隔，默认为5秒
|
||||
queue_manager = WorkflowQueueManager(monitor_interval=5)
|
||||
345
app/comfy/comfy_run.py
Normal file
345
app/comfy/comfy_run.py
Normal file
@@ -0,0 +1,345 @@
|
||||
import json
|
||||
import logging
|
||||
import random
|
||||
import uuid
|
||||
from typing import Optional
|
||||
|
||||
import aiohttp
|
||||
import websockets
|
||||
from aiohttp import ClientTimeout
|
||||
|
||||
from app.comfy.comfy_workflow import ComfyWorkflow, API_OUTPUT_PREFIX
|
||||
from app.comfy.comfy_server import ComfyUIServerInfo, server_manager
|
||||
from app.database.api import (
|
||||
create_workflow_run,
|
||||
create_workflow_run_nodes,
|
||||
get_workflow_run,
|
||||
update_workflow_run_status,
|
||||
update_workflow_run_node_status,
|
||||
)
|
||||
|
||||
|
||||
class ComfyUIExecutionError(Exception):
|
||||
"""ComfyUI执行错误"""
|
||||
|
||||
def __init__(self, error_data):
|
||||
self.error_data = error_data
|
||||
super().__init__(f"ComfyUI execution error: {error_data}")
|
||||
|
||||
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class ComfyRun:
|
||||
"""
|
||||
ComfyUI工作流运行实例，封装单个工作流运行的所有操作
|
||||
"""
|
||||
|
||||
def __init__(self, workflow: ComfyWorkflow, run_id: str, request_data: dict):
|
||||
"""
|
||||
初始化工作流运行实例
|
||||
|
||||
Args:
|
||||
workflow: ComfyWorkflow实例
|
||||
run_id: 运行ID
|
||||
request_data: 请求数据
|
||||
"""
|
||||
self.workflow = workflow
|
||||
self.run_id = run_id
|
||||
self.request_data = request_data
|
||||
self.client_id = str(uuid.uuid4())
|
||||
self.prompt_id: Optional[str] = None
|
||||
self.server: Optional[ComfyUIServerInfo] = None
|
||||
|
||||
@classmethod
|
||||
async def create(cls, workflow: ComfyWorkflow, request_data: dict) -> "ComfyRun":
|
||||
"""
|
||||
创建新的ComfyRun实例并保存到数据库
|
||||
|
||||
Args:
|
||||
workflow: ComfyWorkflow实例
|
||||
request_data: 请求数据
|
||||
|
||||
Returns:
|
||||
ComfyRun实例
|
||||
"""
|
||||
run_id = str(uuid.uuid4())
|
||||
|
||||
# 创建任务记录
|
||||
await create_workflow_run(
|
||||
workflow_run_id=run_id,
|
||||
workflow_name=workflow.workflow_name,
|
||||
workflow_json=json.dumps(workflow.workflow_data.model_dump()),
|
||||
api_spec=json.dumps(workflow.get_api_spec().model_dump()),
|
||||
request_data=json.dumps(request_data),
|
||||
)
|
||||
|
||||
# 创建工作流节点记录
|
||||
nodes_data = []
|
||||
for node in workflow.workflow_data.nodes:
|
||||
nodes_data.append({"id": str(node.id), "type": node.type})
|
||||
await create_workflow_run_nodes(run_id, nodes_data)
|
||||
|
||||
return cls(workflow, run_id, request_data)
|
||||
|
||||
@classmethod
|
||||
async def from_run_id(cls, run_id: str) -> Optional["ComfyRun"]:
|
||||
"""
|
||||
从run_id创建ComfyRun实例
|
||||
|
||||
Args:
|
||||
run_id: 运行ID
|
||||
|
||||
Returns:
|
||||
ComfyRun实例，如果找不到则返回None
|
||||
"""
|
||||
workflow_run = await get_workflow_run(run_id)
|
||||
if not workflow_run:
|
||||
return None
|
||||
|
||||
workflow_data = json.loads(workflow_run.workflow_json)
|
||||
workflow = ComfyWorkflow(workflow_run.workflow_name, workflow_data)
|
||||
request_data = json.loads(workflow_run.request_data)
|
||||
|
||||
return cls(workflow, run_id, request_data)
|
||||
|
||||
async def execute(self, server: ComfyUIServerInfo) -> dict:
|
||||
"""
|
||||
在指定的服务器上执行工作流
|
||||
|
||||
Args:
|
||||
server: ComfyUI服务器信息
|
||||
|
||||
Returns:
|
||||
执行结果
|
||||
"""
|
||||
self.server = server
|
||||
|
||||
try:
|
||||
# 分配服务器资源
|
||||
await server_manager.allocate_server(server.name)
|
||||
|
||||
# 构建prompt
|
||||
prompt = await self.workflow.build_prompt(server, self.request_data)
|
||||
|
||||
# 更新运行状态为running
|
||||
await self._update_status(
|
||||
"running", server.http_url, client_id=self.client_id
|
||||
)
|
||||
|
||||
# 提交到ComfyUI
|
||||
self.prompt_id = await self._queue_prompt(prompt, server.http_url)
|
||||
|
||||
# 更新prompt_id
|
||||
await self._update_status(
|
||||
"running", server.http_url, self.prompt_id, self.client_id
|
||||
)
|
||||
|
||||
logger.info(
|
||||
f"工作流 {self.run_id} 已在 {server.http_url} 上入队，Prompt ID: {self.prompt_id}"
|
||||
)
|
||||
|
||||
# 获取执行结果
|
||||
results = await self._get_execution_results(server.ws_url)
|
||||
|
||||
# 标记完成
|
||||
await self._update_status(
|
||||
"completed", result=json.dumps(results, ensure_ascii=False)
|
||||
)
|
||||
|
||||
return results
|
||||
|
||||
except Exception as e:
|
||||
# 标记失败
|
||||
await self._update_status("failed", error_message=str(e))
|
||||
raise
|
||||
|
||||
finally:
|
||||
# 释放服务器资源
|
||||
try:
|
||||
await server_manager.release_server(server.name)
|
||||
except Exception as e:
|
||||
logger.error(f"释放服务器资源时出错: {e}")
|
||||
|
||||
async def _update_status(
|
||||
self,
|
||||
status: str,
|
||||
server_url: Optional[str] = None,
|
||||
prompt_id: Optional[str] = None,
|
||||
client_id: Optional[str] = None,
|
||||
result: Optional[str] = None,
|
||||
error_message: Optional[str] = None,
|
||||
):
|
||||
"""更新工作流运行状态"""
|
||||
await update_workflow_run_status(
|
||||
self.run_id, status, server_url, prompt_id, client_id, error_message, result
|
||||
)
|
||||
|
||||
async def _queue_prompt(self, prompt: dict, http_url: str) -> str:
|
||||
"""提交工作流到ComfyUI服务器"""
|
||||
# 添加随机缓存破坏器避免缓存
|
||||
for node_id in prompt:
|
||||
prompt[node_id]["inputs"][
|
||||
f"cache_buster_{uuid.uuid4().hex}"
|
||||
] = random.random()
|
||||
|
||||
payload = {
|
||||
"prompt": prompt,
|
||||
"client_id": self.client_id,
|
||||
"extra_data": {
|
||||
"api_key_comfy_org": "",
|
||||
"extra_pnginfo": {"workflow": self.workflow.workflow_data.model_dump()},
|
||||
},
|
||||
}
|
||||
|
||||
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
|
||||
|
||||
async def _get_execution_results(self, ws_url: str) -> dict:
|
||||
"""
|
||||
通过WebSocket连接获取执行结果，支持节点级别的状态跟踪
|
||||
"""
|
||||
aggregated_outputs = {}
|
||||
full_ws_url = f"{ws_url}?clientId={self.client_id}"
|
||||
|
||||
try:
|
||||
async with websockets.connect(full_ws_url) as websocket:
|
||||
logger.info(f"已连接到WebSocket: {full_ws_url}")
|
||||
|
||||
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") == self.prompt_id):
|
||||
continue
|
||||
|
||||
# 捕获并处理执行错误
|
||||
if msg_type == "execution_error":
|
||||
error_data = data
|
||||
logger.error(
|
||||
f"ComfyUI执行错误 (Prompt ID: {self.prompt_id}): {error_data}"
|
||||
)
|
||||
|
||||
# 更新节点状态为失败
|
||||
node_id = error_data.get("node_id")
|
||||
if node_id:
|
||||
await update_workflow_run_node_status(
|
||||
self.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: {self.prompt_id})"
|
||||
)
|
||||
|
||||
# 更新节点状态为运行中
|
||||
await update_workflow_run_node_status(
|
||||
self.run_id, node_id, "running"
|
||||
)
|
||||
|
||||
# 处理节点执行完成
|
||||
elif msg_type == "executed":
|
||||
await self._handle_node_executed(data, aggregated_outputs)
|
||||
|
||||
# 处理整个工作流执行完成
|
||||
elif msg_type == "executing" and data.get("node") is None:
|
||||
logger.info(f"Prompt ID: {self.prompt_id} 执行完成。")
|
||||
return aggregated_outputs
|
||||
|
||||
except websockets.exceptions.ConnectionClosed as e:
|
||||
logger.warning(
|
||||
f"WebSocket 连接已关闭 (Prompt ID: {self.prompt_id})。错误: {e}"
|
||||
)
|
||||
return aggregated_outputs
|
||||
except Exception as e:
|
||||
# 重新抛出我们自己的异常，或者处理其他意外错误
|
||||
if not isinstance(e, ComfyUIExecutionError):
|
||||
logger.error(
|
||||
f"处理 prompt {self.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
|
||||
except Exception as e:
|
||||
logger.error(f"WebSocket连接出错: {e}")
|
||||
raise
|
||||
|
||||
async def _handle_node_executed(self, data: dict, aggregated_outputs: dict):
|
||||
"""处理节点执行完成事件"""
|
||||
node_id = data.get("node")
|
||||
output_data = data.get("output")
|
||||
|
||||
if not node_id or not output_data:
|
||||
return
|
||||
|
||||
# 查找对应的节点
|
||||
node = next(
|
||||
(x for x in self.workflow.workflow_data.nodes if str(x.id) == node_id),
|
||||
None,
|
||||
)
|
||||
|
||||
if not node:
|
||||
return
|
||||
|
||||
# 如果是输出节点，收集结果
|
||||
if node.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: {self.prompt_id})")
|
||||
|
||||
# 更新节点状态为完成
|
||||
await update_workflow_run_node_status(
|
||||
self.run_id,
|
||||
node_id,
|
||||
"completed",
|
||||
output_data=json.dumps(output_data),
|
||||
)
|
||||
|
||||
def get_status_info(self) -> dict:
|
||||
"""获取运行状态信息"""
|
||||
return {
|
||||
"run_id": self.run_id,
|
||||
"workflow_name": self.workflow.workflow_name,
|
||||
"client_id": self.client_id,
|
||||
"prompt_id": self.prompt_id,
|
||||
"server_url": self.server.http_url if self.server else None,
|
||||
}
|
||||
447
app/comfy/comfy_server.py
Normal file
447
app/comfy/comfy_server.py
Normal file
@@ -0,0 +1,447 @@
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import time
|
||||
from datetime import datetime, timedelta
|
||||
from typing import Dict, List, Optional, Set
|
||||
from dataclasses import dataclass, asdict
|
||||
from enum import Enum
|
||||
|
||||
import aiohttp
|
||||
from aiohttp import ClientTimeout
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
from sqlalchemy import select, update, delete
|
||||
from sqlalchemy.orm import selectinload
|
||||
|
||||
from app.database.connection import AsyncSessionLocal
|
||||
from app.database.models import ComfyUIServer as ComfyUIServerModel
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class ServerStatus(Enum):
|
||||
"""服务器状态枚举"""
|
||||
|
||||
ONLINE = "online"
|
||||
OFFLINE = "offline"
|
||||
BUSY = "busy"
|
||||
ERROR = "error"
|
||||
|
||||
|
||||
@dataclass
|
||||
class ComfyUIServerInfo:
|
||||
"""ComfyUI服务器信息"""
|
||||
|
||||
name: str # 服务器名称，由注册方自行拟定
|
||||
http_url: str
|
||||
ws_url: str
|
||||
status: ServerStatus = ServerStatus.OFFLINE
|
||||
last_health_check: Optional[datetime] = None
|
||||
current_tasks: int = 0
|
||||
max_concurrent_tasks: int = 1
|
||||
capabilities: Dict[str, any] = None # 服务器能力信息，如支持的模型等
|
||||
metadata: Dict[str, any] = None # 其他元数据
|
||||
|
||||
def __post_init__(self):
|
||||
if self.capabilities is None:
|
||||
self.capabilities = {}
|
||||
if self.metadata is None:
|
||||
self.metadata = {}
|
||||
|
||||
@classmethod
|
||||
def from_model(cls, model: ComfyUIServerModel) -> "ComfyUIServerInfo":
|
||||
"""从数据库模型创建实例"""
|
||||
capabilities = json.loads(model.capabilities) if model.capabilities else {}
|
||||
metadata = json.loads(model.server_metadata) if model.server_metadata else {}
|
||||
|
||||
return cls(
|
||||
name=model.name,
|
||||
http_url=model.http_url,
|
||||
ws_url=model.ws_url,
|
||||
status=ServerStatus(model.status),
|
||||
last_health_check=model.last_health_check,
|
||||
current_tasks=model.current_tasks,
|
||||
max_concurrent_tasks=model.max_concurrent_tasks,
|
||||
capabilities=capabilities,
|
||||
metadata=metadata,
|
||||
)
|
||||
|
||||
def to_model(self) -> ComfyUIServerModel:
|
||||
"""转换为数据库模型"""
|
||||
return ComfyUIServerModel(
|
||||
name=self.name,
|
||||
http_url=self.http_url,
|
||||
ws_url=self.ws_url,
|
||||
status=self.status.value,
|
||||
last_health_check=self.last_health_check,
|
||||
current_tasks=self.current_tasks,
|
||||
max_concurrent_tasks=self.max_concurrent_tasks,
|
||||
capabilities=json.dumps(self.capabilities),
|
||||
server_metadata=json.dumps(self.metadata),
|
||||
)
|
||||
|
||||
|
||||
class ComfyUIServerManager:
|
||||
"""ComfyUI服务器管理器"""
|
||||
|
||||
def __init__(self):
|
||||
self.servers: Dict[str, ComfyUIServerInfo] = {} # name -> server_info
|
||||
self.lock = asyncio.Lock()
|
||||
self.health_check_interval = 30 # 健康检查间隔（秒）
|
||||
self._health_check_task: Optional[asyncio.Task] = None
|
||||
self._initialized = False
|
||||
|
||||
async def initialize(self):
|
||||
"""初始化管理器，从数据库加载服务器信息"""
|
||||
if self._initialized:
|
||||
return
|
||||
|
||||
async with self.lock:
|
||||
try:
|
||||
await self._load_servers_from_db()
|
||||
self._initialized = True
|
||||
logger.info(f"从数据库加载了 {len(self.servers)} 个服务器")
|
||||
except Exception as e:
|
||||
logger.error(f"从数据库加载服务器失败: {e}")
|
||||
self._initialized = True # 即使失败也标记为已初始化，避免重复尝试
|
||||
|
||||
async def _load_servers_from_db(self):
|
||||
"""从数据库加载服务器信息"""
|
||||
async with AsyncSessionLocal() as session:
|
||||
async with session.begin():
|
||||
result = await session.execute(select(ComfyUIServerModel))
|
||||
models = result.scalars().all()
|
||||
|
||||
for model in models:
|
||||
try:
|
||||
server_info = ComfyUIServerInfo.from_model(model)
|
||||
self.servers[server_info.name] = server_info
|
||||
except Exception as e:
|
||||
logger.error(f"加载服务器 {model.name} 失败: {e}")
|
||||
|
||||
async def _save_server_to_db(self, server: ComfyUIServerInfo):
|
||||
"""保存服务器信息到数据库"""
|
||||
try:
|
||||
async with AsyncSessionLocal() as session:
|
||||
async with session.begin():
|
||||
# 检查是否已存在
|
||||
existing = await session.execute(
|
||||
select(ComfyUIServerModel).where(
|
||||
ComfyUIServerModel.name == server.name
|
||||
)
|
||||
)
|
||||
existing_model = existing.scalar_one_or_none()
|
||||
|
||||
if existing_model:
|
||||
# 更新现有记录
|
||||
await session.execute(
|
||||
update(ComfyUIServerModel)
|
||||
.where(ComfyUIServerModel.name == server.name)
|
||||
.values(
|
||||
http_url=server.http_url,
|
||||
ws_url=server.ws_url,
|
||||
status=server.status.value,
|
||||
last_health_check=server.last_health_check,
|
||||
current_tasks=server.current_tasks,
|
||||
max_concurrent_tasks=server.max_concurrent_tasks,
|
||||
capabilities=json.dumps(server.capabilities),
|
||||
server_metadata=json.dumps(server.metadata),
|
||||
updated_at=datetime.utcnow(),
|
||||
)
|
||||
)
|
||||
else:
|
||||
# 创建新记录
|
||||
new_model = server.to_model()
|
||||
session.add(new_model)
|
||||
|
||||
await session.commit()
|
||||
except Exception as e:
|
||||
logger.error(f"保存服务器 {server.name} 到数据库失败: {e}")
|
||||
|
||||
async def _delete_server_from_db(self, name: str):
|
||||
"""从数据库删除服务器"""
|
||||
try:
|
||||
async with AsyncSessionLocal() as session:
|
||||
async with session.begin():
|
||||
await session.execute(
|
||||
delete(ComfyUIServerModel).where(
|
||||
ComfyUIServerModel.name == name
|
||||
)
|
||||
)
|
||||
await session.commit()
|
||||
except Exception as e:
|
||||
logger.error(f"从数据库删除服务器 {name} 失败: {e}")
|
||||
|
||||
async def register_server(
|
||||
self,
|
||||
name: str,
|
||||
http_url: str,
|
||||
ws_url: str,
|
||||
max_concurrent_tasks: int = 1,
|
||||
capabilities: Optional[Dict[str, any]] = None,
|
||||
metadata: Optional[Dict[str, any]] = None,
|
||||
) -> bool:
|
||||
"""注册新的ComfyUI服务器"""
|
||||
# 确保已初始化
|
||||
await self.initialize()
|
||||
|
||||
# 确保健康检查任务已启动
|
||||
await self._ensure_health_check_started()
|
||||
|
||||
async with self.lock:
|
||||
# 解析URL获取IP和端口
|
||||
from urllib.parse import urlparse
|
||||
|
||||
parsed_url = urlparse(http_url)
|
||||
host = parsed_url.hostname
|
||||
port = parsed_url.port or (80 if parsed_url.scheme == 'http' else 443)
|
||||
|
||||
# 检查是否已存在相同IP+端口的服务器
|
||||
existing_server_name = None
|
||||
for existing_server in self.servers.values():
|
||||
existing_parsed = urlparse(existing_server.http_url)
|
||||
existing_host = existing_parsed.hostname
|
||||
existing_port = existing_parsed.port or (80 if existing_parsed.scheme == 'http' else 443)
|
||||
|
||||
if existing_host == host and existing_port == port:
|
||||
existing_server_name = existing_server.name
|
||||
logger.info(f"检测到IP {host}:{port} 已被服务器 {existing_server_name} 使用，将自动更新配置")
|
||||
break
|
||||
|
||||
# 如果IP+端口已存在，更新现有服务器
|
||||
if existing_server_name:
|
||||
# 删除旧的服务器记录
|
||||
if existing_server_name in self.servers:
|
||||
del self.servers[existing_server_name]
|
||||
logger.info(f"已删除旧服务器 {existing_server_name} 的记录")
|
||||
|
||||
# 检查名称是否已存在
|
||||
if name in self.servers:
|
||||
logger.warning(f"服务器名称 {name} 已存在，将更新配置")
|
||||
|
||||
server_info = ComfyUIServerInfo(
|
||||
name=name,
|
||||
http_url=http_url.rstrip("/"),
|
||||
ws_url=ws_url.rstrip("/"),
|
||||
status=ServerStatus.OFFLINE, # 注册时状态为离线，等待健康检查
|
||||
max_concurrent_tasks=max_concurrent_tasks,
|
||||
capabilities=capabilities or {},
|
||||
metadata=metadata or {},
|
||||
)
|
||||
|
||||
# 立即进行健康检查
|
||||
if await self._check_server_health(server_info):
|
||||
server_info.status = ServerStatus.ONLINE
|
||||
server_info.last_health_check = datetime.now()
|
||||
|
||||
self.servers[name] = server_info
|
||||
|
||||
# 保存到数据库
|
||||
await self._save_server_to_db(server_info)
|
||||
|
||||
# 如果更新了现有服务器，从数据库删除旧记录
|
||||
if existing_server_name:
|
||||
await self._delete_server_from_db(existing_server_name)
|
||||
|
||||
logger.info(f"服务器 {name} 注册成功: {http_url} (IP: {host}:{port})")
|
||||
return True
|
||||
|
||||
async def unregister_server(self, name: str) -> bool:
|
||||
"""注销ComfyUI服务器"""
|
||||
async with self.lock:
|
||||
if name in self.servers:
|
||||
del self.servers[name]
|
||||
# 从数据库删除
|
||||
await self._delete_server_from_db(name)
|
||||
logger.info(f"服务器 {name} 已注销")
|
||||
return True
|
||||
return False
|
||||
|
||||
async def get_available_server(
|
||||
self, required_capabilities: Optional[Dict[str, any]] = None
|
||||
) -> Optional[ComfyUIServerInfo]:
|
||||
"""获取可用的服务器（负载均衡）"""
|
||||
# 确保已初始化
|
||||
await self.initialize()
|
||||
|
||||
# 确保健康检查任务已启动
|
||||
await self._ensure_health_check_started()
|
||||
|
||||
async with self.lock:
|
||||
available_servers = []
|
||||
|
||||
for server in self.servers.values():
|
||||
if (
|
||||
server.status == ServerStatus.ONLINE
|
||||
and server.current_tasks < server.max_concurrent_tasks
|
||||
):
|
||||
|
||||
# 检查能力要求
|
||||
if required_capabilities:
|
||||
if not self._check_capabilities(
|
||||
server.capabilities, required_capabilities
|
||||
):
|
||||
continue
|
||||
|
||||
available_servers.append(server)
|
||||
|
||||
if not available_servers:
|
||||
return None
|
||||
|
||||
# 简单的负载均衡：选择当前任务数最少的服务器
|
||||
return min(available_servers, key=lambda s: s.current_tasks)
|
||||
|
||||
async def allocate_server(self, name: str) -> bool:
|
||||
"""分配服务器资源"""
|
||||
async with self.lock:
|
||||
if name in self.servers:
|
||||
server = self.servers[name]
|
||||
if (
|
||||
server.status == ServerStatus.ONLINE
|
||||
and server.current_tasks < server.max_concurrent_tasks
|
||||
):
|
||||
server.current_tasks += 1
|
||||
if server.current_tasks >= server.max_concurrent_tasks:
|
||||
server.status = ServerStatus.BUSY
|
||||
# 保存到数据库
|
||||
await self._save_server_to_db(server)
|
||||
return True
|
||||
return False
|
||||
|
||||
async def release_server(self, name: str) -> bool:
|
||||
"""释放服务器资源"""
|
||||
async with self.lock:
|
||||
if name in self.servers:
|
||||
server = self.servers[name]
|
||||
if server.current_tasks > 0:
|
||||
server.current_tasks -= 1
|
||||
if (
|
||||
server.status == ServerStatus.BUSY
|
||||
and server.current_tasks < server.max_concurrent_tasks
|
||||
):
|
||||
server.status = ServerStatus.ONLINE
|
||||
# 保存到数据库
|
||||
await self._save_server_to_db(server)
|
||||
return True
|
||||
return False
|
||||
|
||||
async def get_server_status(self, name: str) -> Optional[ComfyUIServerInfo]:
|
||||
"""获取服务器状态"""
|
||||
# 确保已初始化
|
||||
await self.initialize()
|
||||
|
||||
async with self.lock:
|
||||
return self.servers.get(name)
|
||||
|
||||
async def get_all_servers(self) -> List[ComfyUIServerInfo]:
|
||||
"""获取所有服务器信息"""
|
||||
# 确保已初始化
|
||||
await self.initialize()
|
||||
|
||||
async with self.lock:
|
||||
return list(self.servers.values())
|
||||
|
||||
def _check_capabilities(
|
||||
self, server_capabilities: Dict[str, any], required_capabilities: Dict[str, any]
|
||||
) -> bool:
|
||||
"""检查服务器是否满足能力要求"""
|
||||
for key, value in required_capabilities.items():
|
||||
if key not in server_capabilities:
|
||||
return False
|
||||
if isinstance(value, dict) and isinstance(server_capabilities[key], dict):
|
||||
if not self._check_capabilities(server_capabilities[key], value):
|
||||
return False
|
||||
elif server_capabilities[key] != value:
|
||||
return False
|
||||
return True
|
||||
|
||||
async def _check_server_health(self, server: ComfyUIServerInfo) -> bool:
|
||||
"""检查服务器健康状态"""
|
||||
try:
|
||||
timeout = ClientTimeout(total=10)
|
||||
async with aiohttp.ClientSession(timeout=timeout) as session:
|
||||
async with session.get(f"{server.http_url}/system_stats") as response:
|
||||
if response.status == 200:
|
||||
return True
|
||||
except Exception as e:
|
||||
logger.debug(f"服务器 {server.name} 健康检查失败: {e}")
|
||||
|
||||
return False
|
||||
|
||||
async def _health_check_loop(self):
|
||||
"""健康检查循环"""
|
||||
while True:
|
||||
try:
|
||||
await self._perform_health_checks()
|
||||
except Exception as e:
|
||||
logger.error(f"健康检查过程中发生错误: {e}")
|
||||
|
||||
await asyncio.sleep(self.health_check_interval)
|
||||
|
||||
async def _perform_health_checks(self):
|
||||
"""执行健康检查"""
|
||||
current_time = datetime.now()
|
||||
|
||||
async with self.lock:
|
||||
for server in self.servers.values():
|
||||
# 定期健康检查
|
||||
if (
|
||||
not server.last_health_check
|
||||
or current_time - server.last_health_check
|
||||
> timedelta(seconds=self.health_check_interval)
|
||||
):
|
||||
|
||||
is_healthy = await self._check_server_health(server)
|
||||
server.last_health_check = current_time
|
||||
|
||||
if is_healthy and server.status != ServerStatus.ONLINE:
|
||||
server.status = ServerStatus.ONLINE
|
||||
logger.info(f"服务器 {server.name} 恢复在线")
|
||||
elif not is_healthy and server.status == ServerStatus.ONLINE:
|
||||
server.status = ServerStatus.ERROR
|
||||
logger.warning(f"服务器 {server.name} 健康检查失败")
|
||||
|
||||
# 保存到数据库
|
||||
await self._save_server_to_db(server)
|
||||
|
||||
def _start_health_check(self):
|
||||
"""启动健康检查任务"""
|
||||
try:
|
||||
# 检查是否有运行中的事件循环
|
||||
loop = asyncio.get_running_loop()
|
||||
if self._health_check_task is None or self._health_check_task.done():
|
||||
self._health_check_task = asyncio.create_task(self._health_check_loop())
|
||||
except RuntimeError:
|
||||
# 没有运行中的事件循环，延迟启动
|
||||
logger.debug("没有运行中的事件循环，健康检查将在事件循环启动后自动开始")
|
||||
|
||||
async def _ensure_health_check_started(self):
|
||||
"""确保健康检查任务已启动"""
|
||||
if self._health_check_task is None or self._health_check_task.done():
|
||||
self._start_health_check()
|
||||
|
||||
async def shutdown(self):
|
||||
"""关闭管理器"""
|
||||
if self._health_check_task and not self._health_check_task.done():
|
||||
self._health_check_task.cancel()
|
||||
try:
|
||||
await self._health_check_task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
|
||||
# 全局服务器管理器实例
|
||||
server_manager = ComfyUIServerManager()
|
||||
|
||||
|
||||
# 兼容性函数，用于替代原有的静态配置
|
||||
async def get_available_servers() -> List[ComfyUIServerInfo]:
|
||||
"""获取可用的服务器列表（兼容性函数）"""
|
||||
return await server_manager.get_all_servers()
|
||||
|
||||
|
||||
async def get_server_for_task(
|
||||
required_capabilities: Optional[Dict[str, any]] = None
|
||||
) -> Optional[ComfyUIServerInfo]:
|
||||
"""为任务获取合适的服务器（兼容性函数）"""
|
||||
return await server_manager.get_available_server(required_capabilities)
|
||||
760
app/comfy/comfy_workflow.py
Normal file
760
app/comfy/comfy_workflow.py
Normal file
@@ -0,0 +1,760 @@
|
||||
import logging
|
||||
from typing import Any, Dict, List, Optional, Union
|
||||
|
||||
import aiohttp
|
||||
from pydantic import BaseModel
|
||||
from app.comfy.comfy_server import ComfyUIServerInfo
|
||||
|
||||
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
API_INPUT_PREFIX = "INPUT_"
|
||||
API_OUTPUT_PREFIX = "OUTPUT_"
|
||||
|
||||
|
||||
class ComfyWorkflowNodeWidget(BaseModel):
|
||||
"""节点控件定义"""
|
||||
|
||||
name: str
|
||||
|
||||
|
||||
class ComfyWorkflowNodeInput(BaseModel):
|
||||
"""节点输入定义"""
|
||||
|
||||
label: Optional[str] = None
|
||||
localized_name: Optional[str] = None
|
||||
name: str
|
||||
shape: Optional[int] = None
|
||||
type: str
|
||||
link: Optional[int] = None
|
||||
widget: Optional[ComfyWorkflowNodeWidget] = None
|
||||
|
||||
|
||||
class ComfyWorkflowNodeOutput(BaseModel):
|
||||
"""节点输出定义"""
|
||||
|
||||
label: Optional[str] = None
|
||||
localized_name: Optional[str] = None
|
||||
name: str
|
||||
type: str
|
||||
links: Optional[List[int]] = None
|
||||
|
||||
|
||||
class ComfyWorkflowNodeProperties(BaseModel):
|
||||
"""节点属性定义"""
|
||||
|
||||
cnr_id: Optional[str] = None
|
||||
aux_id: Optional[str] = None
|
||||
ver: Optional[str] = None
|
||||
Node_name_for_SR: Optional[str] = None
|
||||
|
||||
|
||||
class ComfyAPIFieldSpec(BaseModel):
|
||||
"""API字段规范"""
|
||||
|
||||
type: str
|
||||
widget_name: str
|
||||
|
||||
|
||||
class ComfyAPIOutputSpec(BaseModel):
|
||||
"""API输出规范"""
|
||||
|
||||
output_name: str
|
||||
|
||||
|
||||
class ComfyAPINodeSpec(BaseModel):
|
||||
"""API节点规范"""
|
||||
|
||||
node_id: str
|
||||
class_type: str
|
||||
inputs: Optional[Dict[str, ComfyAPIFieldSpec]] = {}
|
||||
outputs: Optional[Dict[str, ComfyAPIOutputSpec]] = {}
|
||||
|
||||
|
||||
class ComfyAPISpec(BaseModel):
|
||||
"""API规范定义"""
|
||||
|
||||
inputs: Dict[str, ComfyAPINodeSpec] = {}
|
||||
outputs: Dict[str, ComfyAPINodeSpec] = {}
|
||||
|
||||
|
||||
class ComfyJSONSchemaProperty(BaseModel):
|
||||
"""JSON Schema属性定义"""
|
||||
|
||||
type: str
|
||||
title: Optional[str] = None
|
||||
description: Optional[str] = None
|
||||
default: Optional[Any] = None
|
||||
minimum: Optional[Union[int, float]] = None
|
||||
maximum: Optional[Union[int, float]] = None
|
||||
enum: Optional[List[str]] = None
|
||||
format: Optional[str] = None
|
||||
contentMediaType: Optional[str] = None
|
||||
pattern: Optional[str] = None
|
||||
|
||||
|
||||
class ComfyJSONSchemaNode(BaseModel):
|
||||
"""JSON Schema节点定义"""
|
||||
|
||||
type: str = "object"
|
||||
title: str
|
||||
description: str
|
||||
properties: Dict[str, ComfyJSONSchemaProperty] = {}
|
||||
required: List[str] = []
|
||||
additionalProperties: bool = False
|
||||
|
||||
|
||||
class ComfyJSONSchema(BaseModel):
|
||||
"""JSON Schema定义"""
|
||||
|
||||
schema: str = "http://json-schema.org/draft-07/schema#"
|
||||
type: str = "object"
|
||||
title: str = "ComfyUI Workflow Input Schema"
|
||||
description: str = "工作流输入参数定义"
|
||||
properties: Dict[str, ComfyJSONSchemaNode] = {}
|
||||
required: List[str] = []
|
||||
|
||||
class Config:
|
||||
fields = {"schema": "$schema"}
|
||||
|
||||
|
||||
class ComfyWorkflowNode(BaseModel):
|
||||
"""工作流节点定义"""
|
||||
|
||||
id: int
|
||||
type: str
|
||||
pos: List[float]
|
||||
size: List[float]
|
||||
flags: Dict[str, Any] = {}
|
||||
order: int
|
||||
mode: int
|
||||
inputs: Optional[List[ComfyWorkflowNodeInput]] = []
|
||||
outputs: Optional[List[ComfyWorkflowNodeOutput]] = []
|
||||
properties: ComfyWorkflowNodeProperties = ComfyWorkflowNodeProperties()
|
||||
widgets_values: Optional[List[Any]] = []
|
||||
title: Optional[str] = None
|
||||
|
||||
|
||||
class ComfyWorkflowExtra(BaseModel):
|
||||
"""工作流额外配置"""
|
||||
|
||||
ds: Optional[Dict[str, Any]] = None
|
||||
frontendVersion: Optional[str] = None
|
||||
VHS_latentpreview: Optional[bool] = None
|
||||
VHS_latentpreviewrate: Optional[int] = None
|
||||
VHS_MetadataImage: Optional[bool] = None
|
||||
VHS_KeepIntermediate: Optional[bool] = None
|
||||
|
||||
|
||||
class ComfyWorkflowDataSpec(BaseModel):
|
||||
"""ComfyUI工作流数据规范"""
|
||||
|
||||
id: str
|
||||
revision: int = 0
|
||||
last_node_id: int
|
||||
last_link_id: int
|
||||
nodes: List[ComfyWorkflowNode]
|
||||
links: List[List[Union[int, str]]] = []
|
||||
groups: List[Any] = []
|
||||
config: Dict[str, Any] = {}
|
||||
extra: ComfyWorkflowExtra = ComfyWorkflowExtra()
|
||||
version: float
|
||||
|
||||
|
||||
class ComfyWorkflow:
|
||||
"""
|
||||
ComfyUI工作流处理器类，提供面向对象的工作流管理和处理能力
|
||||
"""
|
||||
|
||||
def __init__(self, workflow_name: str, workflow_data: dict):
|
||||
"""
|
||||
初始化工作流实例
|
||||
|
||||
Args:
|
||||
workflow_name: 工作流名称
|
||||
workflow_data: 工作流数据
|
||||
"""
|
||||
self.workflow_data = ComfyWorkflowDataSpec.model_validate(workflow_data)
|
||||
self.workflow_name = workflow_name
|
||||
self._nodes_map = {str(node.id): node for node in self.workflow_data.nodes}
|
||||
self._api_spec = self._parse_api_spec()
|
||||
self._inputs_json_schema = self._parse_inputs_json_schema()
|
||||
|
||||
async def build_prompt(
|
||||
self, server: ComfyUIServerInfo, request_data: dict[str, Any]
|
||||
) -> dict:
|
||||
"""
|
||||
构建prompt
|
||||
|
||||
Args:
|
||||
request_data: 请求数据
|
||||
|
||||
Returns:
|
||||
dict: 构建好的prompt
|
||||
"""
|
||||
patched_workflow = await self._patch_workflow(server, request_data)
|
||||
if "nodes" not in patched_workflow:
|
||||
raise ValueError("无效的工作流格式")
|
||||
|
||||
prompt_api_format = {}
|
||||
|
||||
# 建立从link_id到源节点的映射
|
||||
link_map = {}
|
||||
for link_data in patched_workflow.get("links", []):
|
||||
link_id, origin_node_id, origin_slot_index = (
|
||||
link_data[0],
|
||||
link_data[1],
|
||||
link_data[2],
|
||||
)
|
||||
|
||||
# 键是目标节点的输入link_id
|
||||
link_map[link_id] = [str(origin_node_id), origin_slot_index]
|
||||
|
||||
for node in patched_workflow["nodes"]:
|
||||
node_id = str(node["id"])
|
||||
inputs_dict = self._get_node_inputs(node, link_map)
|
||||
prompt_api_format[node_id] = {
|
||||
"inputs": inputs_dict,
|
||||
"class_type": node["type"],
|
||||
"_meta": {
|
||||
"title": node.get("title", ""),
|
||||
},
|
||||
}
|
||||
|
||||
# 过滤掉类型为 Note 的节点
|
||||
prompt_api_format = {
|
||||
k: v
|
||||
for k, v in prompt_api_format.items()
|
||||
if v["class_type"] not in ["Note"]
|
||||
}
|
||||
|
||||
return prompt_api_format
|
||||
|
||||
def get_api_spec(self) -> ComfyAPISpec:
|
||||
"""
|
||||
获取API规范
|
||||
|
||||
Returns:
|
||||
ComfyAPISpec: API规范
|
||||
"""
|
||||
return self._api_spec
|
||||
|
||||
def get_inputs_json_schema(self) -> ComfyJSONSchema:
|
||||
"""
|
||||
获取输入参数的JSON Schema
|
||||
|
||||
Returns:
|
||||
ComfyJSONSchema: JSON Schema
|
||||
"""
|
||||
return self._inputs_json_schema
|
||||
|
||||
def _parse_api_spec(self) -> ComfyAPISpec:
|
||||
"""
|
||||
解析工作流，并根据规范生成嵌套结构的API参数名。
|
||||
结构: {'节点名': {'node_id': '节点ID', '字段名': {字段信息}}}
|
||||
"""
|
||||
inputs = {}
|
||||
outputs = {}
|
||||
|
||||
for node_id, node in self._nodes_map.items():
|
||||
title: str = node.title
|
||||
if not title:
|
||||
continue
|
||||
|
||||
if title.startswith(API_INPUT_PREFIX):
|
||||
# 提取节点名（去掉API_INPUT_PREFIX前缀）
|
||||
node_name = title[len(API_INPUT_PREFIX) :]
|
||||
|
||||
# 如果节点名不在inputs中，创建新的节点条目
|
||||
if node_name not in inputs:
|
||||
inputs[node_name] = ComfyAPINodeSpec(
|
||||
node_id=node_id, class_type=node.type, inputs={}, outputs={}
|
||||
)
|
||||
|
||||
if node.inputs:
|
||||
for a_input in node.inputs:
|
||||
if a_input.link is None and a_input.widget:
|
||||
widget_name = a_input.widget.name
|
||||
|
||||
# 特殊处理：隐藏LoadImage节点的upload字段
|
||||
if node.type == "LoadImage" and widget_name == "upload":
|
||||
continue
|
||||
|
||||
# 直接使用widget_name作为字段名
|
||||
inputs[node_name].inputs[widget_name] = ComfyAPIFieldSpec(
|
||||
type=self._get_param_type(a_input.type or "STRING"),
|
||||
widget_name=a_input.widget.name,
|
||||
)
|
||||
|
||||
elif title.startswith(API_OUTPUT_PREFIX):
|
||||
# 提取节点名（去掉API_OUTPUT_PREFIX前缀）
|
||||
node_name = title[len(API_OUTPUT_PREFIX) :]
|
||||
|
||||
# 如果节点名不在outputs中，创建新的节点条目
|
||||
if node_name not in outputs:
|
||||
outputs[node_name] = ComfyAPINodeSpec(
|
||||
node_id=node_id, class_type=node.type, inputs={}, outputs={}
|
||||
)
|
||||
|
||||
if node.outputs:
|
||||
for an_output in node.outputs:
|
||||
output_name = an_output.name
|
||||
|
||||
# 直接使用output_name作为字段名
|
||||
outputs[node_name].outputs[output_name] = ComfyAPIOutputSpec(
|
||||
output_name=an_output.name
|
||||
)
|
||||
|
||||
return ComfyAPISpec(inputs=inputs, outputs=outputs)
|
||||
|
||||
def _parse_inputs_json_schema(self) -> ComfyJSONSchema:
|
||||
"""
|
||||
解析工作流，生成符合 JSON Schema 标准的 API 规范。
|
||||
返回一个只包含输入参数的 JSON Schema 对象。
|
||||
"""
|
||||
# 收集所有输入节点
|
||||
properties = {}
|
||||
required = []
|
||||
|
||||
for _, node in self._nodes_map.items():
|
||||
title: str = node.title
|
||||
if not title:
|
||||
continue
|
||||
|
||||
if title.startswith(API_INPUT_PREFIX):
|
||||
# 提取节点名（去掉API_INPUT_PREFIX前缀）
|
||||
node_name = title[len(API_INPUT_PREFIX) :]
|
||||
|
||||
node_properties = {}
|
||||
node_required = []
|
||||
|
||||
if node.inputs:
|
||||
for a_input in node.inputs:
|
||||
if a_input.link is None and a_input.widget:
|
||||
widget_name = a_input.widget.name
|
||||
|
||||
# 特殊处理：隐藏LoadImage节点的upload字段
|
||||
if node.type == "LoadImage" and widget_name == "upload":
|
||||
continue
|
||||
|
||||
# 获取节点的当前值作为默认值
|
||||
# 创建空的link_map，因为输入节点通常没有连接
|
||||
empty_link_map = {}
|
||||
node_current_inputs = self._get_node_inputs(
|
||||
node, empty_link_map
|
||||
)
|
||||
current_value = node_current_inputs.get(widget_name)
|
||||
|
||||
# 生成字段的 JSON Schema 定义，包含默认值
|
||||
field_schema = self._generate_field_schema_model(
|
||||
a_input, current_value
|
||||
)
|
||||
node_properties[widget_name] = field_schema
|
||||
|
||||
# 如果字段是必需的，添加到required列表
|
||||
# TODO: 从输入配置中获取是否必需的信息
|
||||
node_required.append(widget_name)
|
||||
|
||||
# 创建节点的JSON Schema
|
||||
if node_properties:
|
||||
properties[node_name] = ComfyJSONSchemaNode(
|
||||
title=f"输入节点: {node_name}",
|
||||
description=f"节点类型: {node.type}",
|
||||
properties=node_properties,
|
||||
required=node_required,
|
||||
)
|
||||
|
||||
# 如果节点有必需字段，将整个节点标记为必需
|
||||
if node_required:
|
||||
required.append(node_name)
|
||||
|
||||
return ComfyJSONSchema(properties=properties, required=required)
|
||||
|
||||
async def _patch_workflow(
|
||||
self,
|
||||
server: ComfyUIServerInfo,
|
||||
request_data: dict[str, Any],
|
||||
) -> dict:
|
||||
"""
|
||||
将request_data中的参数值，patch到workflow_data中。并返回修改后的workflow_data。
|
||||
|
||||
request_data结构: {"节点名称": {"字段名称": "字段值"}}
|
||||
"""
|
||||
|
||||
nodes_map = {k: v.model_copy(deep=True) for k, v in self._nodes_map.items()}
|
||||
|
||||
for node_name, node_fields in request_data.items():
|
||||
if node_name not in self._api_spec.inputs:
|
||||
continue
|
||||
|
||||
node_spec = self._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 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.inputs or []) if inp.widget is not None
|
||||
]
|
||||
|
||||
target_widget_index = -1
|
||||
for i, widget_input in enumerate(widget_inputs):
|
||||
if (
|
||||
widget_input.widget
|
||||
and widget_input.widget.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 not target_node.widgets_values or not isinstance(
|
||||
target_node.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.type == "LoadImage":
|
||||
value = await self._upload_image_to_comfy(server, value)
|
||||
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副本
|
||||
patched_workflow_data = self.workflow_data.model_copy(deep=True)
|
||||
patched_workflow_data.nodes = list(nodes_map.values())
|
||||
|
||||
# 对workflow_data中的nodes进行排序，按照id从小到大
|
||||
patched_workflow_data.nodes = sorted(
|
||||
patched_workflow_data.nodes, key=lambda x: x.id
|
||||
)
|
||||
|
||||
# 转换为字典格式以保持向后兼容
|
||||
return patched_workflow_data.model_dump()
|
||||
|
||||
def _convert_workflow_to_prompt_api_format(self, 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 = (
|
||||
link_data[0],
|
||||
link_data[1],
|
||||
link_data[2],
|
||||
)
|
||||
|
||||
# 键是目标节点的输入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 = self._get_node_inputs(node, link_map)
|
||||
prompt_api_format[node_id] = {
|
||||
"inputs": inputs_dict,
|
||||
"class_type": node["type"],
|
||||
"_meta": {
|
||||
"title": node.get("title", ""),
|
||||
},
|
||||
}
|
||||
|
||||
# 过滤掉类型为 Note 的节点
|
||||
prompt_api_format = {
|
||||
k: v
|
||||
for k, v in prompt_api_format.items()
|
||||
if v["class_type"] not in ["Note"]
|
||||
}
|
||||
|
||||
return prompt_api_format
|
||||
|
||||
async def _upload_image_to_comfy(
|
||||
self, server: ComfyUIServerInfo, file_path: str
|
||||
) -> str:
|
||||
"""
|
||||
上传文件到服务器。
|
||||
file 是一个http链接
|
||||
返回一个名称
|
||||
"""
|
||||
file_name = file_path.split("/")[-1]
|
||||
file_data = await self._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(self, 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(
|
||||
self, input_field: dict, default_value: Any = None
|
||||
) -> dict:
|
||||
"""
|
||||
根据输入字段信息生成 JSON Schema 字段定义（字典版本，向后兼容）
|
||||
|
||||
Args:
|
||||
input_field: 输入字段配置
|
||||
default_value: 字段的默认值
|
||||
"""
|
||||
field_schema = {
|
||||
"title": input_field["widget"]["name"],
|
||||
"description": input_field.get("description", ""),
|
||||
}
|
||||
|
||||
# 根据字段类型设置 schema 类型
|
||||
field_type = self._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",
|
||||
"contentMediaType": "image/*",
|
||||
"pattern": ".*\\.(jpg|jpeg|png)$",
|
||||
"description": "[ext:jpg,jpeg,png][file-size:1MB] 上传图片文件",
|
||||
}
|
||||
)
|
||||
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_value is not None:
|
||||
field_schema["default"] = default_value
|
||||
elif "default" in input_field:
|
||||
field_schema["default"] = input_field["default"]
|
||||
|
||||
return field_schema
|
||||
|
||||
def _generate_field_schema_model(
|
||||
self, input_field: ComfyWorkflowNodeInput, default_value: Any = None
|
||||
) -> ComfyJSONSchemaProperty:
|
||||
"""
|
||||
根据输入字段信息生成 JSON Schema 字段定义（模型版本）
|
||||
|
||||
Args:
|
||||
input_field: 输入字段配置
|
||||
default_value: 字段的默认值
|
||||
"""
|
||||
title = input_field.widget.name if input_field.widget else input_field.name
|
||||
description = input_field.localized_name or ""
|
||||
|
||||
# 根据字段类型设置 schema 类型
|
||||
field_type = self._get_param_type(input_field.type or "STRING")
|
||||
|
||||
if field_type == "int":
|
||||
return ComfyJSONSchemaProperty(
|
||||
type="integer",
|
||||
title=title,
|
||||
description=description,
|
||||
default=default_value,
|
||||
# TODO: 从输入配置中获取min/max值
|
||||
)
|
||||
elif field_type == "float":
|
||||
return ComfyJSONSchemaProperty(
|
||||
type="number",
|
||||
title=title,
|
||||
description=description,
|
||||
default=default_value,
|
||||
# TODO: 从输入配置中获取min/max值
|
||||
)
|
||||
elif field_type == "UploadFile":
|
||||
return ComfyJSONSchemaProperty(
|
||||
type="string",
|
||||
title=title,
|
||||
format="binary",
|
||||
contentMediaType="image/*",
|
||||
pattern=".*\\.(jpg|jpeg|png)$",
|
||||
description="[ext:jpg,jpeg,png][file-size:1MB] 上传图片文件",
|
||||
default=default_value,
|
||||
)
|
||||
else: # string
|
||||
return ComfyJSONSchemaProperty(
|
||||
type="string",
|
||||
title=title,
|
||||
description=description,
|
||||
default=default_value,
|
||||
# TODO: 从输入配置中获取枚举值
|
||||
)
|
||||
|
||||
def _get_param_type(self, 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"
|
||||
|
||||
def _get_node_inputs(
|
||||
self, node: Union[ComfyWorkflowNode, dict], link_map: dict
|
||||
) -> dict:
|
||||
"""
|
||||
获取节点的输入字段值
|
||||
|
||||
Args:
|
||||
node: 节点数据（可以是ComfyWorkflowNode模型或字典）
|
||||
link_map: 从link_id到源节点的映射
|
||||
|
||||
Returns:
|
||||
dict: 包含节点输入字段值的字典
|
||||
"""
|
||||
inputs_dict = {}
|
||||
|
||||
# 兼容处理：如果是字典，转换为模型
|
||||
if isinstance(node, dict):
|
||||
try:
|
||||
node_model = ComfyWorkflowNode.model_validate(node)
|
||||
except Exception:
|
||||
# 如果转换失败，使用原来的字典方式
|
||||
return self._get_node_inputs_dict(node, link_map)
|
||||
else:
|
||||
node_model = node
|
||||
|
||||
# 1. 处理控件输入 (widgets)
|
||||
widgets_values = node_model.widgets_values or []
|
||||
widget_cursor = 0
|
||||
|
||||
for input_config in node_model.inputs or []:
|
||||
# 如果是widget并且有对应的widgets_values
|
||||
if input_config.widget and widget_cursor < len(widgets_values):
|
||||
widget_name = input_config.widget.name
|
||||
if widget_name:
|
||||
# 使用widgets_values中的值，因为这里已经包含了API传入的修改
|
||||
inputs_dict[widget_name] = widgets_values[widget_cursor]
|
||||
widget_cursor += 1
|
||||
|
||||
# 特殊处理：如果节点是 EG_CYQ_JB ，则需要忽略widgets_values的第1个值，这个是用来控制control net的，它不参与赋值
|
||||
if node_model.type == "EG_CYQ_JB" and widget_cursor == 1:
|
||||
widget_cursor += 1
|
||||
|
||||
# 2. 处理连接输入 (links)
|
||||
for input_config in node_model.inputs or []:
|
||||
if 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]
|
||||
|
||||
return inputs_dict
|
||||
|
||||
def _get_node_inputs_dict(self, node: dict, link_map: dict) -> dict:
|
||||
"""
|
||||
获取节点的输入字段值（字典版本，用于向后兼容）
|
||||
|
||||
Args:
|
||||
node: 节点数据字典
|
||||
link_map: 从link_id到源节点的映射
|
||||
|
||||
Returns:
|
||||
dict: 包含节点输入字段值的字典
|
||||
"""
|
||||
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
|
||||
|
||||
# 特殊处理：如果节点是 EG_CYQ_JB ，则需要忽略widgets_values的第1个值，这个是用来控制control net的，它不参与赋值
|
||||
if node["type"] == "EG_CYQ_JB" and widget_cursor == 1:
|
||||
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]
|
||||
|
||||
return inputs_dict
|
||||
Reference in New Issue
Block a user