Update app.py

app.py

@@ -91,15 +91,23 @@ 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):
+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
-    return RATIO_MAP.get(size, RATIO_MAP["16:9"]), clean_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
 
 @app.route('/')
 def index():
@@ -116,17 +124,20 @@ 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)
+        image_size, clean_prompt, use_original = extract_params_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(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
+        if use_original:
+            enhanced_prompt = clean_prompt
+        else:
+            try:
+                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
         
         new_url = f'https://api.siliconflow.cn/v1/{model}/text-to-image'
         new_request_body = {
@@ -175,27 +186,32 @@ def handle_request():
         image_data = {'data': [{'url': image_url}]}
         
         if stream:
-            return stream_response(unique_id, image_data, clean_prompt, enhanced_prompt, image_size, current_timestamp, model, system_fingerprint)
+            return stream_response(unique_id, image_data, clean_prompt, enhanced_prompt, image_size, current_timestamp, model, system_fingerprint, use_original)
         else:
-            return non_stream_response(unique_id, image_data, clean_prompt, enhanced_prompt, image_size, current_timestamp, model, system_fingerprint)
+            return non_stream_response(unique_id, image_data, clean_prompt, enhanced_prompt, image_size, current_timestamp, model, system_fingerprint, use_original)
     except Exception as e:
         logger.error(f"Unexpected error in handle_request: {str(e)}")
         return jsonify({"error": f"Internal Server Error: {str(e)}"}), 500
 
-def stream_response(unique_id, image_data, original_prompt, translated_prompt, size, created, model, system_fingerprint):
+def stream_response(unique_id, image_data, original_prompt, translated_prompt, size, created, model, system_fingerprint, use_original):
     logger.debug("Starting stream response")
-    return Response(stream_with_context(generate_stream(unique_id, image_data, original_prompt, translated_prompt, size, created, model, system_fingerprint)), content_type='text/event-stream')
+    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):
+def generate_stream(unique_id, image_data, original_prompt, translated_prompt, size, created, model, system_fingerprint, use_original):
     chunks = [
         f"原始提示词：\n{original_prompt}\n",
-        f"翻译后的提示词：\n{translated_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({
@@ -229,10 +245,13 @@ def generate_stream(unique_id, image_data, original_prompt, translated_prompt, s
     })
     yield f"data: {final_chunk}\n\n"
 
-def non_stream_response(unique_id, image_data, original_prompt, translated_prompt, size, created, model, system_fingerprint):
-    content = (
-        f"原始提示词：{original_prompt}\n"
-        f"翻译后的提示词：{translated_prompt}\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"