from flask import Flask, request, jsonify, Response
import requests
import json
import time
import random

app = Flask(__name__)

SYSTEM_ASSISTANT = """作为 Stable Diffusion Prompt 提示词专家，您将从关键词中创建提示，通常来自 Danbooru 等数据库。
提示通常描述图像，使用常见词汇，按重要性排列，并用逗号分隔。避免使用"-"或"."，但可以接受空格和自然语言。避免词汇重复。

为了强调关键词，请将其放在括号中以增加其权重。例如，"(flowers)"将'flowers'的权重增加1.1倍，而"(((flowers)))"将其增加1.331倍。使用"(flowers:1.5)"将'flowers'的权重增加1.5倍。只为重要的标签增加权重。

提示包括三个部分：**前缀** （质量标签+风格词+效果器）+ **主题** （图像的主要焦点）+ **场景** （背景、环境）。

*   前缀影响图像质量。像"masterpiece"、"best quality"、"4k"这样的标签可以提高图像的细节。像"illustration"、"lensflare"这样的风格词定义图像的风格。像"bestlighting"、"lensflare"、"depthoffield"这样的效果器会影响光照和深度。

*   主题是图像的主要焦点，如角色或场景。对主题进行详细描述可以确保图像丰富而详细。增加主题的权重以增强其清晰度。对于角色，描述面部、头发、身体、服装、姿势等特征。

*   场景描述环境。没有场景，图像的背景是平淡的，主题显得过大。某些主题本身包含场景（例如建筑物、风景）。像"花草草地"、"阳光"、"河流"这样的环境词可以丰富场景。你的任务是设计图像生成的提示。请按照以下步骤进行操作：

1.  我会发送给您一个图像场景。需要你生成详细的图像描述
2.  图像描述必须是英文，输出为Positive Prompt。

示例：

我发送：二战时期的护士。
您回复只回复：
A WWII-era nurse in a German uniform, holding a wine bottle and stethoscope, sitting at a table in white attire, with a table in the background, masterpiece, best quality, 4k, illustration style, best lighting, depth of field, detailed character, detailed environment.
"""

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)}"

def translate_and_enhance_prompt(prompt, auth_token):
    translate_url = 'https://api.siliconflow.cn/v1/chat/completions'
    translate_body = {
        'model': 'Qwen/Qwen2-72B-Instruct',
        'messages': [
            {'role': 'system', 'content': SYSTEM_ASSISTANT},
            {'role': 'user', 'content': prompt}
        ]
    }
    headers = {
        'Content-Type': 'application/json',
        'Authorization': auth_token
    }
    response = requests.post(translate_url, headers=headers, json=translate_body)
    if response.status_code != 200:
        raise Exception(f'Translation error: {response.status_code}')
    result = response.json()
    return result['choices'][0]['message']['content']

@app.route('/')
def index():
    return "text-to-image with siliconflow", 200

@app.route('/ai/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']
        
        # 获取随机token
        random_token = get_random_token(request.headers.get('Authorization'))
        if not random_token:
            return jsonify({"error": "Unauthorized: Invalid or missing Authorization header"}), 401
        
        # 翻译并增强prompt
        enhanced_prompt = translate_and_enhance_prompt(prompt, random_token)
        
        new_url = f'https://api.siliconflow.cn/v1/{model}/text-to-image'
        new_request_body = {
            "prompt": enhanced_prompt,
            "image_size": "1024x1024",
            "batch_size": 1,
            "num_inference_steps": 4,
            "guidance_scale": 1
        }
        
        headers = {
            'accept': 'application/json',
            'content-type': 'application/json',
            '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,
                "object": "chat.completion.chunk",
                "created": current_timestamp,
                "model": model,
                "choices": [
                    {
                        "index": 0,
                        "delta": {
                            "content": f"原始提示词：{prompt}\n增强后的提示词：{enhanced_prompt}\n生成的图像：\n![]({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"原始提示词：{prompt}\n增强后的提示词：{enhanced_prompt}\n生成的图像：\n![]({image_url})"
                        },
                        "logprobs": None,
                        "finish_reason": "length"
                    }
                ],
                "usage": {
                    "prompt_tokens": len(prompt),
                    "completion_tokens": len(enhanced_prompt) + len(image_url),
                    "total_tokens": len(prompt) + len(enhanced_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)