Update app.py

app.py

@@ -6,16 +6,6 @@ import random
 
 app = Flask(__name__)
 
-def get_random_token(auth_header):
-    if not auth_header:
-        return None
-    if auth_header.startswith('Bearer '):
-        auth_header = auth_header[7:]
-    tokens = [token.strip() for token in auth_header.split(',') if token.strip()]
-    if not tokens:
-        return None
-    return f"Bearer {random.choice(tokens)}"
-
 @app.route('/')
 def index():
     return "text-to-image with siliconflow", 200
@@ -27,11 +17,13 @@ def handle_request():
         model = body.get('model')
         messages = body.get('messages')
         stream = body.get('stream', False)
+
         if not model or not messages or len(messages) == 0:
             return jsonify({"error": "Bad Request: Missing required fields"}), 400
-        
+
         prompt = messages[-1]['content']
         new_url = f'https://api.siliconflow.cn/v1/{model}/text-to-image'
+
         new_request_body = {
             "prompt": prompt,
             "image_size": "1024x1024",
@@ -39,24 +31,34 @@ def handle_request():
             "num_inference_steps": 4,
             "guidance_scale": 1
         }
-        
-        # 获取随机token
-        random_token = get_random_token(request.headers.get('Authorization'))
-        if not random_token:
-            return jsonify({"error": "Unauthorized: Invalid or missing Authorization header"}), 401
-        
+
+        # 从传入的 Authorization 头中随机选择一个 token
+        authorization_header = request.headers.get('Authorization')
+        if authorization_header:
+            # 去掉 "Bearer " 前缀并分割 token
+            tokens = authorization_header.replace("Bearer ", "").split(',')
+            if len(tokens) > 1:
+                selected_token = random.choice(tokens).strip()
+            else:
+                selected_token = tokens[0].strip()
+            # 重新格式化为 "Bearer 随机选择的token"
+            selected_token = f"Bearer {selected_token}"
+        else:
+            return jsonify({"error": "Unauthorized: Missing Authorization header"}), 401
+
         headers = {
             'accept': 'application/json',
             'content-type': 'application/json',
-            'Authorization': random_token
+            'Authorization': selected_token
         }
-        
+
         response = requests.post(new_url, headers=headers, json=new_request_body)
         response_body = response.json()
+
         image_url = response_body['images'][0]['url']
         unique_id = int(time.time() * 1000)
         current_timestamp = unique_id // 1000
-        
+
         if stream:
             response_payload = {
                 "id": unique_id,
@@ -100,8 +102,9 @@ def handle_request():
             }
             data_string = json.dumps(response_payload)
             return Response(f"{data_string}\n\n", content_type='text/event-stream')
+
     except Exception as e:
         return jsonify({"error": f"Internal Server Error: {str(e)}"}), 500
 
 if __name__ == '__main__':
-    app.run(host='0.0.0.0', port=8000)
+    app.run(host='0.0.0.0', port=8000)