flux2api / app.py
smgc's picture
Create app.py
52f607b verified
raw
history blame
3.75 kB
from flask import Flask, request, jsonify, Response
import requests
import json
import time
import random
app = Flask(__name__)
@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:
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",
"batch_size": 1,
"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}"
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)
}
}
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)