From cfde3ab545a4cf96a32d18417c703f10515e3f0d Mon Sep 17 00:00:00 2001
From: hwzhuhao <923196325@qq.com>
Date: Sun, 27 Oct 2024 15:13:05 +0800
Subject: [PATCH] feat: add stable-diffusion-3-5-large for the text-to-image
 tool with siliconflow

---
 .../builtin/siliconflow/tools/stable_diffusion.py  | 14 +++++++++-----
 .../siliconflow/tools/stable_diffusion.yaml        |  3 +++
 2 files changed, 12 insertions(+), 5 deletions(-)

diff --git a/api/core/tools/provider/builtin/siliconflow/tools/stable_diffusion.py b/api/core/tools/provider/builtin/siliconflow/tools/stable_diffusion.py
index d6a0b03d1b185..db43790c06aaa 100644
--- a/api/core/tools/provider/builtin/siliconflow/tools/stable_diffusion.py
+++ b/api/core/tools/provider/builtin/siliconflow/tools/stable_diffusion.py
@@ -5,9 +5,12 @@
 from core.tools.entities.tool_entities import ToolInvokeMessage
 from core.tools.tool.builtin_tool import BuiltinTool
 
-SDURL = {
-    "sd_3": "https://api.siliconflow.cn/v1/stabilityai/stable-diffusion-3-medium/text-to-image",
-    "sd_xl": "https://api.siliconflow.cn/v1/stabilityai/stable-diffusion-xl-base-1.0/text-to-image",
+SILICONFLOW_API_URL = "https://api.siliconflow.cn/v1/image/generations"
+
+SD_MODELS = {
+    "sd_3": "stabilityai/stable-diffusion-3-medium",
+    "sd_xl": "stabilityai/stable-diffusion-xl-base-1.0",
+    "sd_3.5_large": "stabilityai/stable-diffusion-3-5-large",
 }
 
 
@@ -22,9 +25,10 @@ def _invoke(
         }
 
         model = tool_parameters.get("model", "sd_3")
-        url = SDURL.get(model)
+        sd_model = SD_MODELS.get(model)
 
         payload = {
+            "model": sd_model,
             "prompt": tool_parameters.get("prompt"),
             "negative_prompt": tool_parameters.get("negative_prompt", ""),
             "image_size": tool_parameters.get("image_size", "1024x1024"),
@@ -34,7 +38,7 @@ def _invoke(
             "num_inference_steps": tool_parameters.get("num_inference_steps", 20),
         }
 
-        response = requests.post(url, json=payload, headers=headers)
+        response = requests.post(SILICONFLOW_API_URL, json=payload, headers=headers)
         if response.status_code != 200:
             return self.create_text_message(f"Got Error Response:{response.text}")
 
diff --git a/api/core/tools/provider/builtin/siliconflow/tools/stable_diffusion.yaml b/api/core/tools/provider/builtin/siliconflow/tools/stable_diffusion.yaml
index dce10adc87f22..b330c92e163a3 100644
--- a/api/core/tools/provider/builtin/siliconflow/tools/stable_diffusion.yaml
+++ b/api/core/tools/provider/builtin/siliconflow/tools/stable_diffusion.yaml
@@ -40,6 +40,9 @@ parameters:
       - value: sd_xl
         label:
           en_US: Stable Diffusion XL
+      - value: sd_3.5_large
+        label:
+          en_US: Stable Diffusion 3.5 Large
     default: sd_3
     label:
       en_US: Choose Image Model