Update app.py

app.py

@@ -56,6 +56,41 @@ RATIO_MAP = {
     "9:16": "1152x2048"
 }
 
+# 模型映射
+MODEL_MAPPING = {
+    "flux.1-schnell": {
+        "provider": "black-forest-labs",
+        "mapping": "black-forest-labs/FLUX.1-schnell"
+    },
+    "sd-turbo": {
+        "provider": "stabilityai",
+        "mapping": "stabilityai/sd-turbo"
+    },
+    "sdxl-turbo": {
+        "provider": "stabilityai",
+        "mapping": "stabilityai/sdxl-turbo"
+    },
+    "stable-diffusion-2-1": {
+        "provider": "stabilityai",
+        "mapping": "stabilityai/stable-diffusion-2-1"
+    },
+    "stable-diffusion-3-medium": {
+        "provider": "stabilityai",
+        "mapping": "stabilityai/stable-diffusion-3-medium"
+    },
+    "stable-diffusion-xl-base-1.0": {
+        "provider": "stabilityai",
+        "mapping": "stabilityai/stable-diffusion-xl-base-1.0"
+    }
+}
+
+# 模拟身份验证函数
+def getAuthCookie(req):
+    auth_cookie = req.headers.get('Authorization')
+    if auth_cookie and auth_cookie.startswith('Bearer '):
+        return auth_cookie
+    return None
+
 def get_random_token(auth_header):
     if not auth_header:
         return None
@@ -149,6 +184,40 @@ def index():
     """
     return usage, 200
 
+@app.route('/ai/v1/models', methods=['GET'])
+def get_models():
+    try:
+        # 验证身份
+        auth_cookie = getAuthCookie(request)
+        if not auth_cookie:
+            return jsonify({"error": "Unauthorized"}), 401
+        
+        # 返回模型列表
+        models_list = [
+            {
+                "id": model_id,
+                "object": "model",
+                "created": int(time.time()),
+                "owned_by": info["provider"],
+                "permission": [],
+                "root": model_id,
+                "parent": None
+            }
+            for model_id, info in MODEL_MAPPING.items()
+        ]
+        
+        # 记录日志
+        app.logger.info(f'{request.remote_addr} - GET /ai/v1/models - 200')
+        
+        return jsonify({
+            "object": "list",
+            "data": models_list
+        })
+    
+    except Exception as error:
+        app.logger.error(f"Error: {str(error)}")
+        return jsonify({"error": "Authentication failed", "details": str(error)}), 401
+
 @app.route('/ai/v1/chat/completions', methods=['POST'])
 def handle_request():
     try: