From 9074a1c480cd912b4fbf1792a8616133d33aab63 Mon Sep 17 00:00:00 2001 From: Eric Wang Date: Sat, 17 May 2025 17:23:34 +0800 Subject: [PATCH] Add API documentation support, register Swagger UI and API routes, add new API documentation /docs, and integrate into main application startup logic. --- .gitignore | 6 + api_server/apispec.py | 27 +++ api_server/routes/__init__.py | 4 + api_server/routes/api_docs.py | 281 +++++++++++++++++++++++++ api_server/routes/internal/__init__.py | 4 + api_server/utils/__init__.py | 4 + api_server/utils/schemas.py | 135 ++++++++++++ main.py | 5 + requirements.txt | 2 + 9 files changed, 468 insertions(+) create mode 100644 api_server/apispec.py create mode 100644 api_server/routes/api_docs.py create mode 100644 api_server/utils/__init__.py create mode 100644 api_server/utils/schemas.py diff --git a/.gitignore b/.gitignore index 4e8cea71e..6ba7b8733 100644 --- a/.gitignore +++ b/.gitignore @@ -13,6 +13,10 @@ extra_model_paths.yaml .idea/ venv/ .venv/ +memory_bank/ +mcp_server_o/ +custom_modes/ +.cursor/ /web/extensions/* !/web/extensions/logging.js.example !/web/extensions/core/ @@ -24,3 +28,5 @@ web_custom_versions/ openapi.yaml filtered-openapi.yaml uv.lock +.gitignore +lab/ diff --git a/api_server/apispec.py b/api_server/apispec.py new file mode 100644 index 000000000..8ec5ac253 --- /dev/null +++ b/api_server/apispec.py @@ -0,0 +1,27 @@ +from aiohttp_apispec import setup_aiohttp_apispec +from api_server.routes.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" + } + ) \ No newline at end of file diff --git a/api_server/routes/__init__.py b/api_server/routes/__init__.py index e69de29bb..3546a5e70 100644 --- a/api_server/routes/__init__.py +++ b/api_server/routes/__init__.py @@ -0,0 +1,4 @@ +# Import API docs related modules +from api_server.routes import api_docs + +__all__ = ['api_docs'] diff --git a/api_server/routes/api_docs.py b/api_server/routes/api_docs.py new file mode 100644 index 000000000..3b11e31d3 --- /dev/null +++ b/api_server/routes/api_docs.py @@ -0,0 +1,281 @@ +from aiohttp_apispec import docs, response_schema, request_schema, querystring_schema +from aiohttp import web +import asyncio +from api_server.utils.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/ \ No newline at end of file diff --git a/api_server/routes/internal/__init__.py b/api_server/routes/internal/__init__.py index e69de29bb..a5ff26b8e 100644 --- a/api_server/routes/internal/__init__.py +++ b/api_server/routes/internal/__init__.py @@ -0,0 +1,4 @@ +# Import internal API route module +from api_server.routes.internal.internal_routes import InternalRoutes + +__all__ = ['InternalRoutes'] diff --git a/api_server/utils/__init__.py b/api_server/utils/__init__.py new file mode 100644 index 000000000..e5571cec3 --- /dev/null +++ b/api_server/utils/__init__.py @@ -0,0 +1,4 @@ +# Import utility modules +from api_server.utils import schemas, file_operations + +__all__ = ['schemas', 'file_operations'] \ No newline at end of file diff --git a/api_server/utils/schemas.py b/api_server/utils/schemas.py new file mode 100644 index 000000000..f3afb825a --- /dev/null +++ b/api_server/utils/schemas.py @@ -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) \ No newline at end of file diff --git a/main.py b/main.py index 0fde6d221..70377acc7 100644 --- a/main.py +++ b/main.py @@ -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() @@ -271,6 +272,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=(q, prompt_server,)).start() if args.quick_test_for_ci: @@ -290,6 +294,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 diff --git a/requirements.txt b/requirements.txt index 8f7a78984..ff9abe6ce 100644 --- a/requirements.txt +++ b/requirements.txt @@ -24,3 +24,5 @@ spandrel soundfile av>=14.2.0 pydantic~=2.0 +aiohttp-apispec +marshmallow