Update app.py

app.py

@@ -2,6 +2,7 @@ from flask import Flask, request, Response
 import requests
 import json
 import time
+from openai import OpenAI
 
 app = Flask(__name__)
 
@@ -35,9 +36,23 @@ def req_flux(token, prompt_value, aspect_ratio="1:1", output_format="webp", num_
         print(f"Error making request: {e}")
         return None
 
-def generate_fake_stream(image_url):
+def generate_optimized_prompt(api_key, api_base, system_prompt, user_input):
+    client = OpenAI(api_key=api_key, base_url=api_base)
+
+    completion = client.chat.completions.create(
+        model="gpt-4o-mini",
+        messages=[
+            {"role": "system", "content": system_prompt},
+            {"role": "user", "content": user_input}
+        ]
+    )
+    return completion.choices[0].message.content
+
+def generate_fake_stream(image_url, optimized_prompt):
     chunks = [
-        "I'm generating an image based on your prompt...",
+        "I'm generating an optimized prompt based on your input...",
+        f"Optimized prompt: {optimized_prompt}",
+        "Now generating an image based on the optimized prompt...",
         "The image is being processed...",
         "Almost there...",
         f"Image generated successfully! Here it is:\n\n![Generated Image]({image_url})"
@@ -56,21 +71,31 @@ def chat_completions():
     stream = data.get('stream', False)
     
     # Extract the prompt from the last user message
-    prompt = next((msg['content'] for msg in reversed(messages) if msg['role'] == 'user'), None)
+    user_input = next((msg['content'] for msg in reversed(messages) if msg['role'] == 'user'), None)
+    
+    if not user_input:
+        return Response(json.dumps({'error': 'No valid user input provided'}), status=400, mimetype='application/json')
     
-    if not prompt:
-        return Response(json.dumps({'error': 'No valid prompt provided'}), status=400, mimetype='application/json')
+    # Generate optimized prompt using GPT-4-mini
+    '''
+    secret_variable：
+    api_key = ""  # Replace with your actual API key
+    api_base = ""  # Replace with your actual API base URL
+    '''
+    system_prompt = """作为 Stable Diffusion Prompt 提示词专家，您将从关键词中创建提示，通常来自 Danbooru 等数据库。提示通常描述图像，使用常见词汇，按重要性排列，并用逗号分隔。避免使用"-"或"."，但可以接受空格和自然语言。避免词汇重复。为了强调关键词，请将其放在括号中以增加其权重。例如，"(flowers)"将'flowers'的权重增加1.1倍，而"(((flowers)))"将其增加1.331倍。使用"(flowers:1.5)"将'flowers'的权重增加1.5倍。只为重要的标签增加权重。提示包括三个部分：前缀（质量标签+风格词+效果器）+ 主题（图像的主要焦点）+ 场景（背景、环境）。前缀影响图像质量。像"masterpiece"、"best quality"、"ultra-detailed"、"high resolution"、"photorealistic" 这样的标签可以显著提高图像的细节和整体质量。像"illustration"、"lensflare"、"cinematic lighting" 这样的风格词定义图像的风格和光影效果。像"best lighting"、"volumetric lighting"、"depth of field" 这样的效果器会影响光照和深度。主题是图像的主要焦点，如角色或场景。对主题进行详细描述可以确保图像丰富而详细。增加主题的权重以增强其清晰度。对于角色，描述面部、头发、身体、服装、姿势等特征，同时加入细致的纹理和高光处理。场景描述环境。没有场景，图像的背景是平淡的，主题显得过大。某些主题本身包含场景（例如建筑物、风景）。像"lush greenery"、"golden sunlight"、"crystal clear river" 这样的环境词可以丰富场景，并增强其视觉吸引力。考虑添加天气效果，如"soft morning mist"、"sunset glow" 来进一步增强场景的氛围。你的任务是设计图像生成的提示。请按照以下步骤进行操作：我会发送给您一个图像场景。需要你生成详细的图像描述。图像描述必须是英文，输出为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, ultra-detailed, high resolution, photorealistic, illustration style, best lighting, volumetric lighting, depth of field, sharp focus, detailed character, richly textured environment."""
+    optimized_prompt = generate_optimized_prompt(api_key, api_base, system_prompt, user_input)
     
+    # Generate image using the optimized prompt
     token = get_token()
     if not token:
         return Response(json.dumps({'error': 'Failed to get token'}), status=500, mimetype='application/json')
     
-    image_url = req_flux(token, prompt)
+    image_url = req_flux(token, optimized_prompt)
     if not image_url:
         return Response(json.dumps({'error': 'Failed to generate image'}), status=500, mimetype='application/json')
     
     if stream:
-        return Response(generate_fake_stream(image_url), mimetype='text/event-stream')
+        return Response(generate_fake_stream(image_url, optimized_prompt), mimetype='text/event-stream')
     else:
         response = {
             "id": f"chatcmpl-{int(time.time())}",
@@ -82,15 +107,15 @@ def chat_completions():
                     "index": 0,
                     "message": {
                         "role": "assistant",
-                        "content": f"I've generated an image based on your prompt. Here it is:\n\n![Generated Image]({image_url})"
+                        "content": f"I've generated an optimized prompt based on your input: '{optimized_prompt}'\n\nUsing this prompt, I've created an image. Here it is:\n\n![Generated Image]({image_url})"
                     },
                     "finish_reason": "stop"
                 }
             ],
             "usage": {
-                "prompt_tokens": len(prompt.split()),
-                "completion_tokens": 20,  # Approximate
-                "total_tokens": len(prompt.split()) + 20
+                "prompt_tokens": len(user_input.split()),
+                "completion_tokens": len(optimized_prompt.split()) + 20,  # Approximate
+                "total_tokens": len(user_input.split()) + len(optimized_prompt.split()) + 20
             }
         }
         return Response(json.dumps(response), mimetype='application/json')