AI生图

main.py

@@ -1,6 +1,7 @@
 from fastapi import FastAPI, Request
 from fastapi.staticfiles import StaticFiles
 from fastapi.middleware.cors import CORSMiddleware
+from fastapi.openapi.utils import get_openapi
 from starlette.middleware.sessions import SessionMiddleware
 from contextlib import asynccontextmanager
 import logging
@@ -41,6 +42,7 @@ async def lifespan(app: FastAPI):
     
     # 确保静态文件目录存在
     os.makedirs(settings.static_dir, exist_ok=True)
+    os.makedirs(settings.media_dir, exist_ok=True)
     
     yield
     
@@ -57,19 +59,55 @@ app = FastAPI(
     description="用于管理墨水屏设备、内容和待办事项的API",
     version="1.0.0",
     lifespan=lifespan,
-    openapi_components={
-        "securitySchemes": {
-            "APIKeyHeader": {
-                "type": "apiKey",
-                "in": "header",
-                "name": "X-API-Key",
-                "description": "API Key鉴权，请在下方输入正确的API Key"
-            }
-        }
-    },
-    security=[{"APIKeyHeader": []}]
 )
 
+# 自定义OpenAPI模式以显示API Key鉴权按钮
+def custom_openapi():
+    if app.openapi_schema:
+        return app.openapi_schema
+    
+    openapi_schema = get_openapi(
+        title=app.title,
+        version=app.version,
+        description=app.description,
+        routes=app.routes,
+    )
+    
+    # 添加安全方案
+    if "components" not in openapi_schema:
+        openapi_schema["components"] = {}
+    
+    security_scheme = {
+        "type": "apiKey",
+        "in": "header",
+        "name": "X-API-Key",
+        "description": "API Key鉴权，请在下方输入正确的API Key"
+    }
+    
+    if "securitySchemes" not in openapi_schema["components"]:
+        openapi_schema["components"]["securitySchemes"] = {}
+        
+    openapi_schema["components"]["securitySchemes"]["APIKeyHeader"] = security_scheme
+    
+    # 添加全局安全要求
+    if "security" not in openapi_schema:
+        openapi_schema["security"] = []
+    
+    # 避免重复添加
+    has_apikey_security = False
+    for security_req in openapi_schema["security"]:
+        if "APIKeyHeader" in security_req:
+            has_apikey_security = True
+            break
+            
+    if not has_apikey_security:
+        openapi_schema["security"].append({"APIKeyHeader": []})
+    
+    app.openapi_schema = openapi_schema
+    return app.openapi_schema
+
+app.openapi = custom_openapi
+
 # 添加CORS中间件
 app.add_middleware(
     CORSMiddleware,
@@ -90,6 +128,7 @@ app.add_middleware(SessionMiddleware, secret_key=settings.secret_key)
 
 # 挂载静态文件
 app.mount("/static", StaticFiles(directory=settings.static_dir), name="static")
+app.mount("/media", StaticFiles(directory=settings.media_dir), name="media")
 
 # 注册API路由
 app.include_router(api_router, prefix="/api")