Update app.py

app.py

@@ -3,13 +3,60 @@ import requests
 import json
 import time
 import random
-import logging
 
 app = Flask(__name__)
 
-# 配置日志格式
-logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(message)s')
-logger = logging.getLogger(__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():
@@ -22,45 +69,40 @@ 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
-
-        authorization_header = request.headers.get('Authorization')
-        if not authorization_header:
-            return jsonify({"error": "Unauthorized: Missing Authorization header"}), 401
-
-        # Extract tokens from Authorization header
-        tokens = authorization_header.split(' ')[1].split(',')
-        if len(tokens) == 1:
-            selected_token = tokens[0]
-        else:
-            selected_token = random.choice(tokens)
-
+        
         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": prompt,
+            "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': f'Bearer {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 = str(int(time.time() * 1000))  # Convert id to string
-        current_timestamp = int(unique_id) // 1000
-
+        unique_id = int(time.time() * 1000)
+        current_timestamp = unique_id // 1000
+        
         if stream:
             response_payload = {
                 "id": unique_id,
@@ -71,7 +113,7 @@ def handle_request():
                     {
                         "index": 0,
                         "delta": {
-                            "content": f"![]({image_url})"
+                            "content": f"原始提示词：{prompt}\n增强后的提示词：{enhanced_prompt}\n生成的图像：\n![]({image_url})"
                         },
                         "finish_reason": "stop"
                     }
@@ -90,7 +132,7 @@ def handle_request():
                         "index": 0,
                         "message": {
                             "role": "assistant",
-                            "content": f"![]({image_url})"
+                            "content": f"原始提示词：{prompt}\n增强后的提示词：{enhanced_prompt}\n生成的图像：\n![]({image_url})"
                         },
                         "logprobs": None,
                         "finish_reason": "length"
@@ -98,19 +140,14 @@ def handle_request():
                 ],
                 "usage": {
                     "prompt_tokens": len(prompt),
-                    "completion_tokens": len(image_url),
-                    "total_tokens": len(prompt) + len(image_url)
+                    "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
 
-    finally:
-        # 记录请求的 model 和 被命中的 token
-        logger.info(f'"POST /ai/v1/chat/completions HTTP/1.1" "model: {model}" "token: {selected_token}" "status: {response.status_code}" -')
-
 if __name__ == '__main__':
-    app.run(host='0.0.0.0', port=8000)
+    app.run(host='0.0.0.0', port=8000)