ComfyUI-WorkflowPublisher/workflow_service/routes/service.py

151 lines
4.3 KiB
Python

import asyncio
import os
from datetime import datetime
from typing import List
import aiohttp
from fastapi import APIRouter, HTTPException, Path
from pydantic import BaseModel
from workflow_service.comfy.comfy_queue import queue_manager
from workflow_service.comfy.comfy_server import server_manager
from workflow_service.config import Settings
settings = Settings()
service_router = APIRouter(
prefix="/api/service",
tags=["Status"],
)
class ServerQueueDetails(BaseModel):
running_count: int
pending_count: int
class ServerStatus(BaseModel):
server_index: int
http_url: str
ws_url: str
is_reachable: bool
is_free: bool
queue_details: ServerQueueDetails
class FileDetails(BaseModel):
name: str
size_kb: float
modified_at: datetime
class ServerFiles(BaseModel):
server_index: int
http_url: str
input_files: List[FileDetails]
output_files: List[FileDetails]
@service_router.get("/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)}")
@service_router.get("/servers_status", response_model=List[ServerStatus])
async def get_servers_status():
"""
获取所有已配置的ComfyUI服务器的配置信息和实时状态。
"""
servers = await server_manager.get_all_servers()
if not servers:
return []
async with aiohttp.ClientSession() as session:
status_tasks = [
queue_manager.get_server_status(server, session) for server in servers
]
live_statuses = await asyncio.gather(*status_tasks)
response_list = []
for i, server_config in enumerate(servers):
status_data = live_statuses[i]
response_list.append(
ServerStatus(
server_index=i,
http_url=server_config.http_url,
ws_url=server_config.ws_url,
is_reachable=status_data["is_reachable"],
is_free=status_data["is_free"],
queue_details=status_data["queue_details"],
)
)
return response_list
@service_router.get("/servers/{server_index}/files", response_model=ServerFiles)
async def list_server_files(
server_index: int = Path(..., ge=0, description="服务器在配置列表中的索引")
):
"""
获取指定ComfyUI服务器的输入和输出文件夹中的文件列表。
注意:由于节点管理不再维护目录,此接口返回空列表。
"""
servers = await server_manager.get_all_servers()
if server_index >= len(servers):
raise HTTPException(
status_code=404,
detail=f"服务器索引 {server_index} 超出范围。有效索引为 0 到 {len(servers) - 1}",
)
server_config = servers[server_index]
# 由于节点管理不再维护目录,返回空列表
input_files, output_files = [], []
return ServerFiles(
server_index=server_index,
http_url=server_config.http_url,
input_files=input_files,
output_files=output_files,
)
async def _get_folder_contents(path: str) -> List[FileDetails]:
"""异步地列出并返回文件夹内容的详细信息。"""
if not os.path.isdir(path):
return []
def sync_list_files(dir_path):
files = []
try:
for entry in os.scandir(dir_path):
if entry.is_file():
stat = entry.stat()
files.append(
FileDetails(
name=entry.name,
size_kb=round(stat.st_size / 1024, 2),
modified_at=datetime.fromtimestamp(stat.st_mtime),
)
)
except OSError as e:
print(f"无法扫描目录 {dir_path}: {e}")
return sorted(files, key=lambda x: x.modified_at, reverse=True)
return await asyncio.to_thread(sync_list_files, path)