Update app.py

app.py

@@ -2,13 +2,14 @@ from flask import Flask, request, jsonify, Response
 import requests
 import json
 import time
+import random
 
 app = Flask(__name__)
 
 @app.route('/')
 def index():
     return "text-to-image with siliconflow", 200
-    
+
 @app.route('/ai/v1/chat/completions', methods=['POST'])
 def handle_request():
     try:
@@ -22,7 +23,6 @@ def handle_request():
 
         prompt = messages[-1]['content']
         new_url = f'https://api.siliconflow.cn/v1/{model}/text-to-image'
-
         new_request_body = {
             "prompt": prompt,
             "image_size": "1024x1024",
@@ -31,16 +31,29 @@ def handle_request():
             "guidance_scale": 1
         }
 
+        # Get the Authorization header
+        auth_header = request.headers.get('Authorization', '')
+
+        # Split the tokens and randomly select one
+        tokens = auth_header.split(',')
+        if tokens:
+            # Remove 'Bearer ' prefix if present
+            tokens = [token.strip().replace('Bearer ', '') for token in tokens]
+            selected_token = random.choice(tokens)
+            auth_header = f'Bearer {selected_token}'
+        else:
+            auth_header = ''
+
         headers = {
             'accept': 'application/json',
             'content-type': 'application/json',
-            'Authorization': request.headers.get('Authorization')
+            'Authorization': auth_header
         }
 
         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
 
@@ -87,7 +100,6 @@ 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