Update app.py

app.py

@@ -3,24 +3,22 @@ import requests
 import json
 import time
 import random
-import logging
 
 app = Flask(__name__)
 
-# 配置日志
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
-
-# 添加 StreamHandler 以实时输出日志到控制台
-console_handler = logging.StreamHandler()
-console_handler.setLevel(logging.INFO)
-formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
-console_handler.setFormatter(formatter)
-logger.addHandler(console_handler)
+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 "flux2api with siliconflow", 200
+    return "text-to-image with siliconflow", 200
 
 @app.route('/ai/v1/chat/completions', methods=['POST'])
 def handle_request():
@@ -29,13 +27,11 @@ 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",
@@ -43,36 +39,24 @@ def handle_request():
             "num_inference_steps": 4,
             "guidance_scale": 1
         }
-
-        # 从传入的 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}"
-            # 记录命中的 token
-            logger.info(f"Request hit token: {selected_token}")
-        else:
-            return jsonify({"error": "Unauthorized: Missing Authorization header"}), 401
-
+        
+        # 获取随机token
+        random_token = get_random_token(request.headers.get('Authorization'))
+        if not random_token:
+            return jsonify({"error": "Unauthorized: Invalid or missing Authorization header"}), 401
+        
         headers = {
             'accept': 'application/json',
             'content-type': 'application/json',
-            'Authorization': selected_token
+            'Authorization': random_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,
@@ -116,9 +100,8 @@ 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)