From d1f8a98ad0057ed3a05bcb4d88897cfd9e71ac59 Mon Sep 17 00:00:00 2001
From: yinpeng <2291314224@qq.com>
Date: Tue, 11 Feb 2025 06:10:55 +0800
Subject: [PATCH] =?UTF-8?q?feat:=20=E6=94=AF=E6=8C=81=E5=9C=A8=E5=9B=BE?=
 =?UTF-8?q?=E7=89=87=E7=94=9F=E6=88=90=E6=8F=90=E7=A4=BA=E8=AF=8D=E4=B8=AD?=
 =?UTF-8?q?=E9=80=9A=E8=BF=87=E6=A0=87=E8=AE=B0=E6=8E=A7=E5=88=B6=E5=8F=82?=
 =?UTF-8?q?=E6=95=B0?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 README.md                            | 25 ++++++++++++++-
 app/services/image_create_service.py | 46 ++++++++++++++++++++++++++++
 2 files changed, 70 insertions(+), 1 deletion(-)

diff --git a/README.md b/README.md
index 181deea..370b1f1 100644
--- a/README.md
+++ b/README.md
@@ -247,12 +247,34 @@ uvicorn app.main:app --host 0.0.0.0 --port 8000 --reload
     ```json
     {
     "model": "dall-e-3",
-    "prompt": "汉服美女",
+    "prompt": "{n:2} {ratio:16:9} 汉服美女",
     "n": 1,
     "size": "1024x1024"
     }
     ```
 
+    **Prompt参数说明:**
+
+    prompt支持通过特殊标记来控制生成参数：
+
+    1. 图片数量控制:
+       - 格式: `{n:数量}`
+       - 示例: `{n:2} 一只可爱的猫` - 生成2张图片
+       - 取值范围: 1-4
+       - 说明: 如果在prompt中指定了n，将覆盖请求body中的n参数
+
+    2. 图片比例控制:
+       - 格式: `{ratio:宽:高}`
+       - 示例: `{ratio:16:9} 一片森林` - 生成16:9比例的图片
+       - 支持的比例: "1:1"、"3:4"、"4:3"、"9:16"、"16:9"
+       - 说明: 如果指定了size参数，将优先使用size对应的比例
+
+    3. 参数组合:
+       - 示例: `{n:2} {ratio:16:9} 一片美丽的森林` - 生成2张16:9比例的图片
+       - 说明: 这些参数标记会自动从prompt中移除，不会影响实际的图片生成提示词
+
+    > 注意：n的取值范围[1,4], ratio取值范围"1:1"、"3:4"、"4:3"、"9:16" 和 "16:9"
+
 ## 📚 代码结构
 
 ```plaintext
@@ -328,6 +350,7 @@ A: 请检查以下几点：
 A: 在请求的 Body 中，将 `stream` 参数设置为 `true` 即可。
 
 **Q: 如何启用代码执行工具？**
+
 A: 在 `.env` 文件的 `TOOLS_CODE_EXECUTION_ENABLED` 变量中, 设置为 `true` 即可。
 
 ## 📄 许可证
diff --git a/app/services/image_create_service.py b/app/services/image_create_service.py
index 4498df0..353c51d 100644
--- a/app/services/image_create_service.py
+++ b/app/services/image_create_service.py
@@ -19,8 +19,42 @@ class ImageCreateService:
         self.paid_key = settings.PAID_KEY
         self.aspect_ratio = aspect_ratio
 
+    def parse_prompt_parameters(self, prompt: str) -> tuple:
+        """从prompt中解析参数
+        支持的格式:
+        - {n:数量} 例如: {n:2} 生成2张图片
+        - {ratio:比例} 例如: {ratio:16:9} 使用16:9比例
+        """
+        import re
+        
+        # 默认值
+        n = 1
+        aspect_ratio = self.aspect_ratio
+        
+        # 解析n参数
+        n_match = re.search(r'{n:(\d+)}', prompt)
+        if n_match:
+            n = int(n_match.group(1))
+            if n < 1 or n > 4:
+                raise ValueError(f"Invalid n value: {n}. Must be between 1 and 4.")
+            prompt = prompt.replace(n_match.group(0), '').strip()
+            
+        # 解析ratio参数    
+        ratio_match = re.search(r'{ratio:(\d+:\d+)}', prompt)
+        if ratio_match:
+            aspect_ratio = ratio_match.group(1)
+            valid_ratios = ["1:1", "3:4", "4:3", "9:16", "16:9"]
+            if aspect_ratio not in valid_ratios:
+                raise ValueError(
+                    f"Invalid ratio: {aspect_ratio}. Must be one of: {', '.join(valid_ratios)}"
+                )
+            prompt = prompt.replace(ratio_match.group(0), '').strip()
+            
+        return prompt, n, aspect_ratio
+
     def generate_images(self, request: ImageGenerationRequest):
         client = genai.Client(api_key=self.paid_key)
+        
         if request.size == "1024x1024":
             self.aspect_ratio = "1:1"
         elif request.size == "1792x1024":
@@ -32,6 +66,18 @@ class ImageCreateService:
                 f"Invalid size: {request.size}. Supported sizes are 1024x1024, 1792x1024, and 1024x1792."
             )
 
+        # 解析prompt中的参数
+        cleaned_prompt, prompt_n, prompt_ratio = self.parse_prompt_parameters(request.prompt)
+        request.prompt = cleaned_prompt
+        
+        # 如果prompt中指定了n，则覆盖请求中的n
+        if prompt_n > 1:
+            request.n = prompt_n
+            
+        # 如果prompt中指定了ratio，则覆盖默认的aspect_ratio
+        if prompt_ratio != self.aspect_ratio:
+            self.aspect_ratio = prompt_ratio
+
         response = client.models.generate_images(
             model=self.image_model,
             prompt=request.prompt,