Update app.py
Browse files
app.py
CHANGED
@@ -56,6 +56,41 @@ RATIO_MAP = {
|
|
56 |
"9:16": "1152x2048"
|
57 |
}
|
58 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
59 |
def get_random_token(auth_header):
|
60 |
if not auth_header:
|
61 |
return None
|
@@ -149,6 +184,40 @@ def index():
|
|
149 |
"""
|
150 |
return usage, 200
|
151 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
152 |
@app.route('/ai/v1/chat/completions', methods=['POST'])
|
153 |
def handle_request():
|
154 |
try:
|
|
|
56 |
"9:16": "1152x2048"
|
57 |
}
|
58 |
|
59 |
+
# 模型映射
|
60 |
+
MODEL_MAPPING = {
|
61 |
+
"flux.1-schnell": {
|
62 |
+
"provider": "black-forest-labs",
|
63 |
+
"mapping": "black-forest-labs/FLUX.1-schnell"
|
64 |
+
},
|
65 |
+
"sd-turbo": {
|
66 |
+
"provider": "stabilityai",
|
67 |
+
"mapping": "stabilityai/sd-turbo"
|
68 |
+
},
|
69 |
+
"sdxl-turbo": {
|
70 |
+
"provider": "stabilityai",
|
71 |
+
"mapping": "stabilityai/sdxl-turbo"
|
72 |
+
},
|
73 |
+
"stable-diffusion-2-1": {
|
74 |
+
"provider": "stabilityai",
|
75 |
+
"mapping": "stabilityai/stable-diffusion-2-1"
|
76 |
+
},
|
77 |
+
"stable-diffusion-3-medium": {
|
78 |
+
"provider": "stabilityai",
|
79 |
+
"mapping": "stabilityai/stable-diffusion-3-medium"
|
80 |
+
},
|
81 |
+
"stable-diffusion-xl-base-1.0": {
|
82 |
+
"provider": "stabilityai",
|
83 |
+
"mapping": "stabilityai/stable-diffusion-xl-base-1.0"
|
84 |
+
}
|
85 |
+
}
|
86 |
+
|
87 |
+
# 模拟身份验证函数
|
88 |
+
def getAuthCookie(req):
|
89 |
+
auth_cookie = req.headers.get('Authorization')
|
90 |
+
if auth_cookie and auth_cookie.startswith('Bearer '):
|
91 |
+
return auth_cookie
|
92 |
+
return None
|
93 |
+
|
94 |
def get_random_token(auth_header):
|
95 |
if not auth_header:
|
96 |
return None
|
|
|
184 |
"""
|
185 |
return usage, 200
|
186 |
|
187 |
+
@app.route('/ai/v1/models', methods=['GET'])
|
188 |
+
def get_models():
|
189 |
+
try:
|
190 |
+
# 验证身份
|
191 |
+
auth_cookie = getAuthCookie(request)
|
192 |
+
if not auth_cookie:
|
193 |
+
return jsonify({"error": "Unauthorized"}), 401
|
194 |
+
|
195 |
+
# 返回模型列表
|
196 |
+
models_list = [
|
197 |
+
{
|
198 |
+
"id": model_id,
|
199 |
+
"object": "model",
|
200 |
+
"created": int(time.time()),
|
201 |
+
"owned_by": info["provider"],
|
202 |
+
"permission": [],
|
203 |
+
"root": model_id,
|
204 |
+
"parent": None
|
205 |
+
}
|
206 |
+
for model_id, info in MODEL_MAPPING.items()
|
207 |
+
]
|
208 |
+
|
209 |
+
# 记录日志
|
210 |
+
app.logger.info(f'{request.remote_addr} - GET /ai/v1/models - 200')
|
211 |
+
|
212 |
+
return jsonify({
|
213 |
+
"object": "list",
|
214 |
+
"data": models_list
|
215 |
+
})
|
216 |
+
|
217 |
+
except Exception as error:
|
218 |
+
app.logger.error(f"Error: {str(error)}")
|
219 |
+
return jsonify({"error": "Authentication failed", "details": str(error)}), 401
|
220 |
+
|
221 |
@app.route('/ai/v1/chat/completions', methods=['POST'])
|
222 |
def handle_request():
|
223 |
try:
|