Merge 26e65a027ee66bca93fc91473c2bda19de85a769 into ecb8d15e7a778ed28e03fc131a8621b2f2945fe8

This commit is contained in:
EricWang 2025-06-12 16:23:18 +08:00 committed by GitHub
commit b835e1b3cf
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
10 changed files with 462 additions and 0 deletions

2
.gitignore vendored
View File

@ -13,6 +13,7 @@ extra_model_paths.yaml
.idea/
venv/
.venv/
.cursor/
/web/extensions/*
!/web/extensions/logging.js.example
!/web/extensions/core/
@ -24,3 +25,4 @@ web_custom_versions/
openapi.yaml
filtered-openapi.yaml
uv.lock

View File

@ -0,0 +1 @@
from .apispec import register_apispec

View File

@ -0,0 +1,281 @@
from aiohttp_apispec import docs, response_schema, request_schema, querystring_schema
from aiohttp import web
import asyncio
from api_server.apispec.schemas import *
def wrap_stable_routes(app):
"""Add Swagger documentation annotations for stable APIs (server.py)"""
from server import PromptServer
server_instance = PromptServer.instance
# ===== GET Methods =====
# GET /prompt - Get queue status
@docs(
tags=["Stable"],
summary="Get current queue status",
description="Return the execution status information of the current queue."
)
@response_schema(QueueStatusSchema(), 200)
async def get_prompt_swagger(request):
return await server_instance.routes._routes_by_method['GET']['/prompt'](request)
app.router.add_get("/prompt", get_prompt_swagger)
# GET /queue - Get queue details
@docs(
tags=["Stable"],
summary="Get queue details",
description="Return detailed information of the current queue."
)
@response_schema(QueueStatusSchema(), 200)
async def get_queue_swagger(request):
return await server_instance.routes._routes_by_method['GET']['/queue'](request)
app.router.add_get("/queue", get_queue_swagger)
# GET /history - Get history
@docs(
tags=["Stable"],
summary="Get history",
description="Return the history of completed generation tasks."
)
@response_schema(HistoryResponseSchema(), 200)
async def get_history_swagger(request):
return await server_instance.routes._routes_by_method['GET']['/history'](request)
app.router.add_get("/history", get_history_swagger)
# GET /history/{prompt_id} - Get specific history
@docs(
tags=["Stable"],
summary="Get specific history",
description="Get detailed information of a specific history by prompt_id."
)
@response_schema(HistoryItemResponseSchema(), 200)
async def get_history_id_swagger(request):
return await server_instance.routes._routes_by_method['GET']['/history/{prompt_id}'](request)
app.router.add_get("/history/{prompt_id}", get_history_id_swagger)
# GET /system_stats - Get system status
@docs(
tags=["Stable"],
summary="Get system status",
description="Return system and resource usage information."
)
@response_schema(SystemStatsSchema(), 200)
async def get_system_stats_swagger(request):
return await server_instance.routes._routes_by_method['GET']['/system_stats'](request)
app.router.add_get("/system_stats", get_system_stats_swagger)
# GET /models - Get model type list
@docs(
tags=["Stable"],
summary="Get model type list",
description="Return the list of supported model types."
)
@response_schema(ModelsListSchema(), 200)
async def list_models_swagger(request):
return await server_instance.routes._routes_by_method['GET']['/models'](request)
app.router.add_get("/models", list_models_swagger)
# GET /models/{folder} - Get model files in a specific type
@docs(
tags=["Stable"],
summary="Get model file list of a specific type",
description="Return the list of model files for the specified type."
)
@response_schema(ModelFilesSchema(), 200)
async def get_models_folder_swagger(request):
return await server_instance.routes._routes_by_method['GET']['/models/{folder}'](request)
app.router.add_get("/models/{folder}", get_models_folder_swagger)
# GET /view - View image
@docs(
tags=["Stable"],
summary="View image",
description="Get and optionally process an image file."
)
@querystring_schema(ViewImageQuerySchema())
async def view_image_swagger(request):
return await server_instance.routes._routes_by_method['GET']['/view'](request)
app.router.add_get("/view", view_image_swagger)
# GET /object_info - Get all node info
@docs(
tags=["Stable"],
summary="Get all node info",
description="Return detailed information of all available nodes."
)
@response_schema(Schema.from_dict({"nodes": fields.Dict(description="Node info dictionary")}), 200)
async def get_object_info_swagger(request):
return await server_instance.routes._routes_by_method['GET']['/object_info'](request)
app.router.add_get("/object_info", get_object_info_swagger)
# GET /object_info/{node_class} - Get specific node info
@docs(
tags=["Stable"],
summary="Get specific node info",
description="Return detailed information of the specified node type."
)
@response_schema(NodeInfoResponseSchema(), 200)
async def get_object_info_node_swagger(request):
return await server_instance.routes._routes_by_method['GET']['/object_info/{node_class}'](request)
app.router.add_get("/object_info/{node_class}", get_object_info_node_swagger)
# ===== POST Methods =====
# POST /prompt - Submit generation task
@docs(
tags=["Stable"],
summary="Submit generation task",
description="Submit a new generation task to the queue."
)
@request_schema(PromptRequestSchema())
@response_schema(PromptResponseSchema(), 200)
async def post_prompt_swagger(request):
return await server_instance.routes._routes_by_method['POST']['/prompt'](request)
app.router.add_post("/prompt", post_prompt_swagger)
# POST /queue - Operate queue
@docs(
tags=["Stable"],
summary="Operate queue",
description="Clear the queue or delete specific queue items."
)
@request_schema(QueueRequestSchema())
async def post_queue_swagger(request):
return await server_instance.routes._routes_by_method['POST']['/queue'](request)
app.router.add_post("/queue", post_queue_swagger)
# POST /history - Operate history
@docs(
tags=["Stable"],
summary="Operate history",
description="Clear the history or delete specific history items."
)
@request_schema(HistoryRequestSchema())
async def post_history_swagger(request):
return await server_instance.routes._routes_by_method['POST']['/history'](request)
app.router.add_post("/history", post_history_swagger)
# POST /interrupt - Interrupt current task
@docs(
tags=["Stable"],
summary="Interrupt current task",
description="Interrupt the currently running generation task."
)
async def post_interrupt_swagger(request):
return await server_instance.routes._routes_by_method['POST']['/interrupt'](request)
app.router.add_post("/interrupt", post_interrupt_swagger)
# POST /free - Release resources
@docs(
tags=["Stable"],
summary="Release resources",
description="Unload models and/or free memory."
)
@request_schema(FreeRequestSchema())
async def post_free_swagger(request):
return await server_instance.routes._routes_by_method['POST']['/free'](request)
app.router.add_post("/free", post_free_swagger)
# POST /upload/image - Upload image
@docs(
tags=["Stable"],
summary="Upload image",
description="Upload an image file to the server."
)
@response_schema(UploadResponseSchema(), 200)
async def upload_image_swagger(request):
return await server_instance.routes._routes_by_method['POST']['/upload/image'](request)
app.router.add_post("/upload/image", upload_image_swagger)
# POST /upload/mask - Upload mask
@docs(
tags=["Stable"],
summary="Upload mask",
description="Upload a mask image and apply it to the original image."
)
@response_schema(UploadResponseSchema(), 200)
async def upload_mask_swagger(request):
return await server_instance.routes._routes_by_method['POST']['/upload/mask'](request)
app.router.add_post("/upload/mask", upload_mask_swagger)
def wrap_internal_routes(app):
"""Add Swagger documentation annotations for internal APIs"""
# GET /internal/logs - Get logs
@docs(
tags=["internal"],
summary="Get logs",
description="Get system log content."
)
@response_schema(LogsResponseSchema(), 200)
async def get_logs_swagger(request):
request_path = web.Request.clone(request)
request_path._match_info = {'tail': 'logs'}
app = request.app.middlewares[0](lambda r: r) # Get parent app
return await app._subapps['/internal']._handle(request_path)
app.router.add_get("/internal/logs", get_logs_swagger)
# GET /internal/logs/raw - Get raw logs
@docs(
tags=["internal"],
summary="Get raw logs",
description="Get raw system logs and terminal size."
)
@response_schema(RawLogsResponseSchema(), 200)
async def get_raw_logs_swagger(request):
request_path = web.Request.clone(request)
request_path._match_info = {'tail': 'logs/raw'}
app = request.app.middlewares[0](lambda r: r) # Get parent app
return await app._subapps['/internal']._handle(request_path)
app.router.add_get("/internal/logs/raw", get_raw_logs_swagger)
# PATCH /internal/logs/subscribe - Subscribe logs
@docs(
tags=["internal"],
summary="Subscribe logs",
description="Enable or disable log subscription for the client."
)
@request_schema(SubscribeLogsRequestSchema())
async def subscribe_logs_swagger(request):
request_path = web.Request.clone(request)
request_path._match_info = {'tail': 'logs/subscribe'}
app = request.app.middlewares[0](lambda r: r) # Get parent app
return await app._subapps['/internal']._handle(request_path)
app.router.add_patch("/internal/logs/subscribe", subscribe_logs_swagger)
# GET /internal/folder_paths - Get folder paths
@docs(
tags=["internal"],
summary="Get folder paths",
description="Get the paths of various types of folders in the system."
)
@response_schema(FolderPathsResponseSchema(), 200)
async def get_folder_paths_swagger(request):
request_path = web.Request.clone(request)
request_path._match_info = {'tail': 'folder_paths'}
app = request.app.middlewares[0](lambda r: r) # Get parent app
return await app._subapps['/internal']._handle(request_path)
app.router.add_get("/internal/folder_paths", get_folder_paths_swagger)
# GET /internal/files/{directory_type} - Get file list
@docs(
tags=["internal"],
summary="Get file list",
description="Get the list of files in the specified type of directory."
)
@response_schema(FilesResponseSchema(), 200)
async def get_files_swagger(request):
directory_type = request.match_info['directory_type']
request_path = web.Request.clone(request)
request_path._match_info = {'tail': f'files/{directory_type}', 'directory_type': directory_type}
app = request.app.middlewares[0](lambda r: r) # Get parent app
return await app._subapps['/internal']._handle(request_path)
app.router.add_get("/internal/files/{directory_type}", get_files_swagger)
def register_api_docs(app):
"""Register API documentation routes on the main app"""
wrap_stable_routes(app) # Stable APIs from server.py
wrap_internal_routes(app) # Internal APIs from api_server/

View File

@ -0,0 +1,27 @@
from aiohttp_apispec import setup_aiohttp_apispec
from api_server.apispec.api_docs import register_api_docs
def register_apispec(app):
"""
Register Swagger UI and API documentation on the given aiohttp app
"""
# Register API documentation
register_api_docs(app)
# Register Swagger UI
setup_aiohttp_apispec(
app=app,
title="ComfyUI API",
version="1.0.0",
url="/docs/swagger.json",
swagger_path="/docs",
swagger_config={
"layout": "StandaloneLayout",
"deepLinking": True,
"displayRequestDuration": True,
"defaultModelsExpandDepth": 3,
"docExpansion": "list",
"tagsSorter": "alpha",
"operationsSorter": "alpha"
}
)

View File

@ -0,0 +1,135 @@
from marshmallow import Schema, fields, validate
# Common Schema
class ErrorResponseSchema(Schema):
"""Common error response model"""
error = fields.Dict(keys=fields.Str(), values=fields.Raw(), required=True, description="Error details")
node_errors = fields.Dict(required=False, description="Node-specific errors")
# Queue-related Schema
class PromptQueueItemSchema(Schema):
"""Queue item model"""
prompt_id = fields.Str(description="Task ID")
number = fields.Float(description="Task number")
class PromptResponseSchema(Schema):
"""Response model for submitting a generation task"""
prompt_id = fields.Str(description="Task ID")
number = fields.Float(description="Task number")
node_errors = fields.Dict(description="Node validation errors")
class QueueStatusSchema(Schema):
"""Queue status response model"""
exec_info = fields.Dict(description="Execution info")
class QueueRequestSchema(Schema):
"""Queue operation request model"""
clear = fields.Bool(description="Whether to clear the queue", required=False)
delete = fields.List(fields.Str(), description="List of queue item IDs to delete", required=False)
# System-related Schema
class SystemStatsSchema(Schema):
"""System status response model"""
system = fields.Dict(description="System info")
devices = fields.List(fields.Dict(), description="Device info")
# Model-related Schema
class ModelsListSchema(Schema):
"""Model type list response"""
model_types = fields.List(fields.Str(), description="List of model types")
class ModelFilesSchema(Schema):
"""Response for model file list in a specific folder"""
files = fields.List(fields.Str(), description="File list")
# Upload-related Schema
class UploadResponseSchema(Schema):
"""Upload response model"""
name = fields.Str(description="File name")
subfolder = fields.Str(description="Subfolder path")
type = fields.Str(description="Upload type")
class UploadImageRequestSchema(Schema):
"""Upload image request model"""
image = fields.Raw(description="Image file", required=True)
overwrite = fields.Bool(description="Whether to overwrite", required=False, default=False)
type = fields.Str(description="Target upload type", required=False, default="input")
subfolder = fields.Str(description="Subfolder path", required=False, default="")
class UploadMaskRequestSchema(Schema):
"""Upload mask request model"""
image = fields.Raw(description="Mask image file", required=True)
original_ref = fields.Dict(description="Original image reference info", required=True)
# View image-related Schema
class ViewImageQuerySchema(Schema):
"""Query parameters for viewing images"""
filename = fields.Str(description="File name", required=True)
type = fields.Str(description="File type", required=False, default="output")
subfolder = fields.Str(description="Subfolder", required=False)
preview = fields.Str(description="Preview parameters", required=False)
channel = fields.Str(description="Channel parameters", required=False, default="rgba")
# History-related Schema
class HistoryResponseSchema(Schema):
"""History response model"""
history = fields.List(fields.Dict(), description="History list")
class HistoryItemResponseSchema(Schema):
"""Single history item response model"""
prompt = fields.Dict(description="Prompt info")
outputs = fields.Dict(description="Output info")
created_at = fields.Str(description="Creation time")
class HistoryRequestSchema(Schema):
"""History operation request model"""
clear = fields.Bool(description="Whether to clear history", required=False)
delete = fields.List(fields.Str(), description="List of history item IDs to delete", required=False)
# Prompt/generation-related Schema
class PromptRequestSchema(Schema):
"""Prompt request model"""
prompt = fields.Dict(description="Workflow graph", required=True)
number = fields.Float(description="Task number", required=False)
front = fields.Bool(description="Insert at the front of the queue", required=False)
extra_data = fields.Dict(description="Extra data", required=False)
client_id = fields.Str(description="Client ID", required=False)
# Node info-related Schema
class NodeInfoResponseSchema(Schema):
"""Node info response model"""
input = fields.Dict(description="Input types")
output = fields.List(fields.Str(), description="Output types")
output_name = fields.List(fields.Str(), description="Output names")
name = fields.Str(description="Node class name")
display_name = fields.Str(description="Display name")
description = fields.Str(description="Node description")
category = fields.Str(description="Node category")
# Internal API Schema
class LogsResponseSchema(Schema):
"""Logs response model"""
logs = fields.Str(description="Log content")
class RawLogsResponseSchema(Schema):
"""Raw logs response model"""
entries = fields.List(fields.Dict(), description="Log entries")
size = fields.Dict(description="Terminal size")
class SubscribeLogsRequestSchema(Schema):
"""Subscribe logs request model"""
clientId = fields.Str(required=True, description="Client ID")
enabled = fields.Bool(required=True, description="Enable subscription")
class FolderPathsResponseSchema(Schema):
"""Folder paths response model"""
paths = fields.Dict(keys=fields.Str(), values=fields.Str(), description="Folder path mapping")
class FilesResponseSchema(Schema):
"""File list response model"""
files = fields.List(fields.Str(), description="File list")
class FreeRequestSchema(Schema):
"""Free memory/model request model"""
unload_models = fields.Bool(description="Whether to unload models", required=False)
free_memory = fields.Bool(description="Whether to free memory", required=False)

View File

@ -0,0 +1 @@

View File

@ -0,0 +1,4 @@
# Import internal API route module
from api_server.routes.internal.internal_routes import InternalRoutes
__all__ = ['InternalRoutes']

View File

@ -0,0 +1,4 @@
# Import utility modules
from api_server.utils import schemas, file_operations
__all__ = ['schemas', 'file_operations']

View File

@ -135,6 +135,7 @@ import comfy.model_management
import comfyui_version
import app.logger
import hook_breaker_ac10a0
from api_server.apispec import register_apispec
def cuda_malloc_warning():
device = comfy.model_management.get_torch_device()
@ -280,6 +281,9 @@ def start_comfyui(asyncio_loop=None):
prompt_server.add_routes()
hijack_progress(prompt_server)
# register Swagger UI to main app
register_apispec(prompt_server.app)
threading.Thread(target=prompt_worker, daemon=True, args=(prompt_server.prompt_queue, prompt_server,)).start()
if args.quick_test_for_ci:
@ -299,6 +303,7 @@ def start_comfyui(asyncio_loop=None):
async def start_all():
await prompt_server.setup()
# start ComfyUI main server
await run(prompt_server, address=args.listen, port=args.port, verbose=not args.dont_print_server, call_on_start=call_on_start)
# Returning these so that other code can integrate with the ComfyUI loop and server

View File

@ -27,3 +27,5 @@ spandrel
soundfile
av>=14.2.0
pydantic~=2.0
aiohttp-apispec
marshmallow<4.0.0