93 lines
2.5 KiB
Python
93 lines
2.5 KiB
Python
#!/usr/bin/env python3
|
|
"""ai-image 入口：REST + Web 控制台 + MCP。
|
|
|
|
启动方式：
|
|
python main.py # Web + REST + MCP(http) 同端口
|
|
python main.py --stdio # 仅 MCP stdio（给 MCP 客户端当子进程用）
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import contextlib
|
|
|
|
from fastapi import FastAPI
|
|
from fastapi.staticfiles import StaticFiles
|
|
|
|
from app import config, db
|
|
from app.mcp_server import build_http_app
|
|
from app.routers import api_gallery, api_generate, api_models, api_providers, api_skills
|
|
|
|
config.ensure_dirs()
|
|
db.init_db()
|
|
|
|
mcp_app = build_http_app(path="/mcp")
|
|
|
|
|
|
@contextlib.asynccontextmanager
|
|
async def lifespan(_: FastAPI):
|
|
async with mcp_app.lifespan(mcp_app):
|
|
yield
|
|
|
|
|
|
app = FastAPI(title="ai-image", version="2.0.0", lifespan=lifespan)
|
|
|
|
app.include_router(api_providers.router)
|
|
app.include_router(api_models.router)
|
|
app.include_router(api_generate.router)
|
|
app.include_router(api_gallery.router)
|
|
app.include_router(api_skills.router)
|
|
|
|
|
|
@app.get("/health")
|
|
def health():
|
|
return {"status": "ok", "output_dir": str(config.OUTPUT_DIR)}
|
|
|
|
|
|
app.mount("/output", StaticFiles(directory=config.OUTPUT_DIR), name="output")
|
|
app.mount("/", StaticFiles(directory=config.WEB_DIR, html=True), name="web")
|
|
|
|
|
|
class RootDispatcher:
|
|
"""按路径前缀把请求分发给 MCP 或 Web 应用。
|
|
|
|
Web 静态资源挂在根路径，MCP 需要独占 /mcp，因此不能靠 mount 处理。
|
|
"""
|
|
|
|
def __init__(self, web_app, mcp_asgi, prefix: str = "/mcp"):
|
|
self.web_app = web_app
|
|
self.mcp_asgi = mcp_asgi
|
|
self.prefix = prefix
|
|
|
|
async def __call__(self, scope, receive, send):
|
|
path = scope.get("path", "")
|
|
if scope["type"] == "http" and (path == self.prefix or path.startswith(self.prefix + "/")):
|
|
await self.mcp_asgi(scope, receive, send)
|
|
else:
|
|
await self.web_app(scope, receive, send)
|
|
|
|
|
|
dispatcher = RootDispatcher(app, mcp_app)
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser(description="ai-image 服务")
|
|
parser.add_argument("--stdio", action="store_true", help="以 MCP stdio 模式运行")
|
|
parser.add_argument("--host", default=config.HOST)
|
|
parser.add_argument("--port", type=int, default=config.PORT)
|
|
args = parser.parse_args()
|
|
|
|
if args.stdio:
|
|
from app.mcp_server import mcp
|
|
|
|
mcp.run()
|
|
return
|
|
|
|
import uvicorn
|
|
|
|
uvicorn.run(dispatcher, host=args.host, port=args.port)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|