Update app.py

app.py

@@ -5,6 +5,7 @@ import time
 import random
 import logging
 import sys
+import re
 
 app = Flask(__name__)
 
@@ -41,6 +42,15 @@ SYSTEM_ASSISTANT = """作为 Stable Diffusion 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 get_random_token(auth_header):
     if not auth_header:
         return None
@@ -81,6 +91,16 @@ def translate_and_enhance_prompt(prompt, auth_token):
         logger.error(f"Error in translate_and_enhance_prompt: {str(e)}")
         raise
 
+def extract_size_from_prompt(prompt):
+    size_match = re.search(r'-s\s+(\S+)', 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
+    return RATIO_MAP.get(size, RATIO_MAP["16:9"]), clean_prompt
+
 @app.route('/')
 def index():
     return "text-to-image with siliconflow", 200
@@ -96,13 +116,14 @@ def handle_request():
             return jsonify({"error": "Bad Request: Missing required fields"}), 400
         
         prompt = messages[-1]['content']
+        image_size, clean_prompt = extract_size_from_prompt(prompt)
         
         random_token = get_random_token(request.headers.get('Authorization'))
         if not random_token:
             return jsonify({"error": "Unauthorized: Invalid or missing Authorization header"}), 401
         
         try:
-            enhanced_prompt = translate_and_enhance_prompt(prompt, random_token)
+            enhanced_prompt = translate_and_enhance_prompt(clean_prompt, random_token)
         except Exception as e:
             logger.error(f"Error in translate_and_enhance_prompt: {str(e)}")
             return jsonify({"error": "Failed to enhance prompt"}), 500
@@ -110,7 +131,7 @@ def handle_request():
         new_url = f'https://api.siliconflow.cn/v1/{model}/text-to-image'
         new_request_body = {
             "prompt": enhanced_prompt,
-            "image_size": "1024x1024",
+            "image_size": image_size,
             "batch_size": 1,
             "num_inference_steps": 4,
             "guidance_scale": 1
@@ -154,9 +175,9 @@ def handle_request():
         image_data = {'data': [{'url': image_url}]}
         
         if stream:
-            return stream_response(unique_id, image_data, prompt, enhanced_prompt, "1024x1024", current_timestamp, model, system_fingerprint)
+            return stream_response(unique_id, image_data, clean_prompt, enhanced_prompt, image_size, current_timestamp, model, system_fingerprint)
         else:
-            return non_stream_response(unique_id, image_data, prompt, enhanced_prompt, "1024x1024", current_timestamp, model, system_fingerprint)
+            return non_stream_response(unique_id, image_data, clean_prompt, enhanced_prompt, image_size, current_timestamp, model, system_fingerprint)
     except Exception as e:
         logger.error(f"Unexpected error in handle_request: {str(e)}")
         return jsonify({"error": f"Internal Server Error: {str(e)}"}), 500