Update main.py

main.py

@@ -360,7 +360,7 @@ def generate_image_stream(unique_id: int, current_timestamp: int, model: str, pr
             {
                 "index": 0,
                 "delta": {
-                    "content": f"\`\`\`\n{{\n  \"prompt\":\"{prompt}\"\n}}\n\`\`\`\n"
+                    "content": f"```\n{{\n  \"prompt\":\"{prompt}\"\n}}\n```\n"
                 }
             }
         ],
@@ -694,21 +694,26 @@ def handle_request():
             if not api_keys or api_keys[0] == "":
                 raise ValueError("API_KEYS environment variable not set")
             
-            # 确保有足够的API密钥
+            # 确保有可用的API密钥
             available_keys = len(api_keys)
             max_images = get_env_int("MAX_IMAGES_PER_REQUEST", 4)
             
             # 用户请求的图片数量不能超过环境变量限制
             requested_count = min(image_count, max_images)
             
-            # 最终生成的图片数量不能超过可用的API密钥数量
-            final_count = min(requested_count, available_keys)
+            # 最终生成的图片数量为用户请求的数量（受环境变量限制）
+            final_count = requested_count
             
             logger.info(f"环境变量限制: {max_images}张, 用户请求: {image_count}张, 可用API密钥: {available_keys}个, 最终生成: {final_count}张")
             
-            # 随机选择不重复的API密钥
-            selected_keys = random.sample(api_keys, final_count)
-            logger.info(f"已选择 {len(selected_keys)} 个API密钥用于图像生成")
+            # 选择API密钥，允许重复使用以实现负载均衡
+            selected_keys = []
+            for i in range(final_count):
+                # 循环使用可用的API密钥
+                key_index = i % available_keys
+                selected_keys.append(api_keys[key_index])
+            
+            logger.info(f"已选择 {len(selected_keys)} 个API密钥用于图像生成（可能包含重复使用的密钥）")
         except ValueError as e:
             logger.error(f"获取外部API密钥失败: {e}")
             return jsonify({"error": "未配置外部API密钥，请设置API_KEYS环境变量"}), 500
@@ -808,7 +813,7 @@ def handle_request():
                         {
                             "index": 0,
                             "delta": {
-                                "content": f"\`\`\`\n{{\n  \"prompt\":\"{safe_prompt}\",\n  \"count\":{final_count}\n}}\n\`\`\`\n"
+                                "content": f"```\n{{\n  \"prompt\":\"{safe_prompt}\",\n  \"count\":{final_count}\n}}\n```\n"
                             }
                         }
                     ],
@@ -1167,3 +1172,5 @@ if __name__ == "__main__":
     
     # 启动服务
     app.run(host="0.0.0.0", port=port)
+
+