Update app.py

app.py

@@ -7,7 +7,6 @@ import logging
 import sys
 import re
 from logging.handlers import TimedRotatingFileHandler
-from datetime import datetime
 
 app = Flask(__name__)
 
@@ -61,6 +60,38 @@ MODEL_MAPPING = {
     }
 }
 
+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.
+"""
+
+RATIO_MAP = {
+    "1:1": "1024x1024",
+    "1:2": "1024x2048",
+    "3:2": "1536x1024",
+    "4:3": "1536x2048",
+    "16:9": "2048x1152",
+    "9:16": "1152x2048"
+}
+
 # 模拟身份验证函数
 def getAuthCookie(req):
     auth_cookie = req.headers.get('Authorization')
@@ -180,8 +211,139 @@ def handle_request():
         app.logger.error(f"Error: {str(e)}")
         return jsonify({"error": f"Internal Server Error: {str(e)}"}), 500
 
-# 其余代码保持不变
-# 例如 stream_response, generate_stream, non_stream_response 等函数
+def extract_params_from_prompt(prompt):
+    size_match = re.search(r'-s\s+(\S+)', prompt)
+    original_match = re.search(r'-o', prompt)
+    
+    if size_match:
+        size = size_match.group(1)
+        clean_prompt = re.sub(r'-s\s+\S+', '', prompt).strip()
+    else:
+        size = "16:9"
+        clean_prompt = prompt
+    
+    use_original = bool(original_match)
+    if use_original:
+        clean_prompt = re.sub(r'-o', '', clean_prompt).strip()
+    
+    image_size = RATIO_MAP.get(size, RATIO_MAP["16:9"])
+    return image_size, clean_prompt, use_original, size
+
+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 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': f'Bearer {auth_token}'
+    }
+    
+    response = requests.post(translate_url, headers=headers, json=translate_body, timeout=30)
+    response.raise_for_status()
+    result = response.json()
+    return result['choices'][0]['message']['content']
+
+def stream_response(unique_id, image_data, original_prompt, translated_prompt, size, created, model, system_fingerprint, use_original):
+    return Response(stream_with_context(generate_stream(unique_id, image_data, original_prompt, translated_prompt, size, created, model, system_fingerprint, use_original)), content_type='text/event-stream')
+
+def generate_stream(unique_id, image_data, original_prompt, translated_prompt, size, created, model, system_fingerprint, use_original):
+    chunks = [
+        f"原始提示词：\n{original_prompt}\n",
+    ]
+    
+    if not use_original:
+        chunks.append(f"翻译后的提示词：\n{translated_prompt}\n")
+    
+    chunks.extend([
+        f"图像规格：{size}\n",
+        "正在根据提示词生成图像...\n",
+        "图像正在处理中...\n",
+        "即将完成...\n",
+        f"生成成功！\n图像生成完毕，以下是结果：\n\n![生成的图像]({image_data['data'][0]['url']})"
+    ])
+
+    for i, chunk in enumerate(chunks):
+        json_chunk = json.dumps({
+            "id": unique_id,
+            "object": "chat.completion.chunk",
+            "created": created,
+            "model": model,
+            "system_fingerprint": system_fingerprint,
+            "choices": [{
+                "index": 0,
+                "delta": {"content": chunk},
+                "logprobs": None,
+                "finish_reason": None
+            }]
+        })
+        yield f"data: {json_chunk}\n\n"
+        time.sleep(0.5)  # 模拟生成时间
+
+    final_chunk = json.dumps({
+        "id": unique_id,
+        "object": "chat.completion.chunk",
+        "created": created,
+        "model": model,
+        "system_fingerprint": system_fingerprint,
+        "choices": [{
+            "index": 0,
+            "delta": {},
+            "logprobs": None,
+            "finish_reason": "stop"
+        }]
+    })
+    yield f"data: {final_chunk}\n\n"
+
+def non_stream_response(unique_id, image_data, original_prompt, translated_prompt, size, created, model, system_fingerprint, use_original):
+    content = f"原始提示词：{original_prompt}\n"
+    
+    if not use_original:
+        content += f"翻译后的提示词：{translated_prompt}\n"
+    
+    content += (
+        f"图像规格：{size}\n"
+        f"图像生成成功！\n"
+        f"以下是结果：\n\n"
+        f"![生成的图像]({image_data['data'][0]['url']})"
+    )
+
+    response = {
+        'id': unique_id,
+        'object': "chat.completion",
+        'created': created,
+        'model': model,
+        'system_fingerprint': system_fingerprint,
+        'choices': [{
+            'index': 0,
+            'message': {
+                'role': "assistant",
+                'content': content
+            },
+            'finish_reason': "stop"
+        }],
+        'usage': {
+            'prompt_tokens': len(original_prompt),
+            'completion_tokens': len(content),
+            'total_tokens': len(original_prompt) + len(content)
+        }
+    }
+
+    return jsonify(response)
 
 if __name__ == '__main__':
     app.run(host='0.0.0.0', port=8000)