Update app.py

app.py

@@ -1,34 +1,29 @@
-from http.server import HTTPServer, BaseHTTPRequestHandler
+from flask import Flask, request, jsonify, Response
+import requests
 import json
-from datetime import datetime
-import urllib.request
-import ssl
-import logging
+import time
+import random
 
-# 设置日志
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
+app = Flask(__name__)
 
-class ChatCompletionsHandler(BaseHTTPRequestHandler):
-    def do_POST(self):
-        if self.path != '/ai/v1/chat/completions':
-            self.send_error(404, "Not Found")
-            return
-
-        content_length = int(self.headers['Content-Length'])
-        post_data = self.rfile.read(content_length)
-        body = json.loads(post_data.decode('utf-8'))
+@app.route('/')
+def index():
+    return "flux2api with siliconflow", 200
 
+@app.route('/v1/chat/completions', methods=['POST'])
+def handle_request():
+    try:
+        body = request.json
         model = body.get('model')
         messages = body.get('messages')
         stream = body.get('stream', False)
 
         if not model or not messages or len(messages) == 0:
-            self.send_error(400, "Bad Request: Missing required fields")
-            return
+            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_url = f'https://api.siliconflow.cn/v1/{model}/text-to-image'
+
         new_request_body = {
             "prompt": prompt,
             "image_size": "1024x1024",
@@ -37,92 +32,79 @@ class ChatCompletionsHandler(BaseHTTPRequestHandler):
             "guidance_scale": 1
         }
 
-        req = urllib.request.Request(new_url, 
-                                     data=json.dumps(new_request_body).encode('utf-8'), 
-                                     headers={
-                                         'accept': 'application/json',
-                                         'content-type': 'application/json',
-                                         'Authorization': self.headers.get('Authorization')
-                                     },
-                                     method='POST')
-
-        ctx = ssl.create_default_context()
-        ctx.check_hostname = False
-        ctx.verify_mode = ssl.CERT_NONE
-
-        try:
-            with urllib.request.urlopen(req, context=ctx) as response:
-                response_body = json.loads(response.read().decode('utf-8'))
-            
-            logger.info(f"API Response: {response_body}")  # 记录 API 响应
-
-            if 'images' not in response_body or not response_body['images']:
-                raise ValueError("No images in the response")
-
-            image_url = response_body['images'][0].get('url')
-            if not image_url:
-                raise ValueError("No URL in the image response")
-
-            unique_id = int(datetime.now().timestamp() * 1000)
-            current_timestamp = int(unique_id / 1000)
-
-            if stream:
-                response_payload = {
-                    "id": unique_id,
-                    "object": "chat.completion.chunk",
-                    "created": current_timestamp,
-                    "model": model,
-                    "choices": [
-                        {
-                            "index": 0,
-                            "delta": {
-                                "content": f"![]({image_url})"
-                            },
-                            "finish_reason": "stop"
-                        }
-                    ]
-                }
-                self.send_response(200)
-                self.send_header('Content-Type', 'text/event-stream')
-                self.end_headers()
-                self.wfile.write(f"data: {json.dumps(response_payload)}\n\n".encode('utf-8'))
+        # 从传入的 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:
-                response_payload = {
-                    "id": unique_id,
-                    "object": "chat.completion",
-                    "created": current_timestamp,
-                    "model": model,
-                    "choices": [
-                        {
-                            "index": 0,
-                            "message": {
-                                "role": "assistant",
-                                "content": f"![]({image_url})"
-                            },
-                            "logprobs": None,
-                            "finish_reason": "length"
-                        }
-                    ],
-                    "usage": {
-                        "prompt_tokens": len(prompt),
-                        "completion_tokens": len(image_url),
-                        "total_tokens": len(prompt) + len(image_url)
+                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': 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,
+                "object": "chat.completion.chunk",
+                "created": current_timestamp,
+                "model": model,
+                "choices": [
+                    {
+                        "index": 0,
+                        "delta": {
+                            "content": f"![]({image_url})"
+                        },
+                        "finish_reason": "stop"
+                    }
+                ]
+            }
+            data_string = json.dumps(response_payload)
+            return Response(f"data: {data_string}\n\n", content_type='text/event-stream')
+        else:
+            response_payload = {
+                "id": unique_id,
+                "object": "chat.completion",
+                "created": current_timestamp,
+                "model": model,
+                "choices": [
+                    {
+                        "index": 0,
+                        "message": {
+                            "role": "assistant",
+                            "content": f"![]({image_url})"
+                        },
+                        "logprobs": None,
+                        "finish_reason": "length"
                     }
+                ],
+                "usage": {
+                    "prompt_tokens": len(prompt),
+                    "completion_tokens": len(image_url),
+                    "total_tokens": len(prompt) + len(image_url)
                 }
-                self.send_response(200)
-                self.send_header('Content-Type', 'application/json')
-                self.end_headers()
-                self.wfile.write(json.dumps(response_payload).encode('utf-8'))
-
-        except Exception as e:
-            logger.error(f"Error occurred: {str(e)}")
-            self.send_error(500, f"Internal Server Error: {str(e)}")
+            }
+            data_string = json.dumps(response_payload)
+            return Response(f"{data_string}\n\n", content_type='text/event-stream')
 
-def run(server_class=HTTPServer, handler_class=ChatCompletionsHandler, port=8000):
-    server_address = ('', port)
-    httpd = server_class(server_address, handler_class)
-    logger.info(f"Starting server on port {port}")
-    httpd.serve_forever()
+    except Exception as e:
+        return jsonify({"error": f"Internal Server Error: {str(e)}"}), 500
 
 if __name__ == '__main__':
-    run()
+    app.run(host='0.0.0.0', port=8000)