From 3d083b758ff7c48614eb30bad0201f3dad8ad351 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E9=9D=9E=E6=B3=95=E6=93=8D=E4=BD=9C?= <hjlarry@163.com>
Date: Sun, 15 Sep 2024 17:14:12 +0800
Subject: [PATCH] feat: add flux dev of siliconflow image-gen tool (#8450)

---
 .../builtin/siliconflow/tools/flux.py         | 11 ++++++++--
 .../builtin/siliconflow/tools/flux.yaml       | 21 ++++++++++++++++---
 2 files changed, 27 insertions(+), 5 deletions(-)

diff --git a/api/core/tools/provider/builtin/siliconflow/tools/flux.py b/api/core/tools/provider/builtin/siliconflow/tools/flux.py
index 1b846624bd1a08..0d16ff385eb30d 100644
--- a/api/core/tools/provider/builtin/siliconflow/tools/flux.py
+++ b/api/core/tools/provider/builtin/siliconflow/tools/flux.py
@@ -5,7 +5,10 @@
 from core.tools.entities.tool_entities import ToolInvokeMessage
 from core.tools.tool.builtin_tool import BuiltinTool
 
-FLUX_URL = "https://api.siliconflow.cn/v1/black-forest-labs/FLUX.1-schnell/text-to-image"
+FLUX_URL = {
+    "schnell": "https://api.siliconflow.cn/v1/black-forest-labs/FLUX.1-schnell/text-to-image",
+    "dev": "https://api.siliconflow.cn/v1/image/generations",
+}
 
 
 class FluxTool(BuiltinTool):
@@ -24,8 +27,12 @@ def _invoke(
             "seed": tool_parameters.get("seed"),
             "num_inference_steps": tool_parameters.get("num_inference_steps", 20),
         }
+        model = tool_parameters.get("model", "schnell")
+        url = FLUX_URL.get(model)
+        if model == "dev":
+            payload["model"] = "black-forest-labs/FLUX.1-dev"
 
-        response = requests.post(FLUX_URL, json=payload, headers=headers)
+        response = requests.post(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/flux.yaml b/api/core/tools/provider/builtin/siliconflow/tools/flux.yaml
index 2a0698700c1577..d06b9bf3e1f489 100644
--- a/api/core/tools/provider/builtin/siliconflow/tools/flux.yaml
+++ b/api/core/tools/provider/builtin/siliconflow/tools/flux.yaml
@@ -6,8 +6,8 @@ identity:
   icon: icon.svg
 description:
   human:
-    en_US: Generate image via SiliconFlow's flux schnell.
-  llm: This tool is used to generate image from prompt via SiliconFlow's flux schnell model.
+    en_US: Generate image via SiliconFlow's flux model.
+  llm: This tool is used to generate image from prompt via SiliconFlow's flux model.
 parameters:
   - name: prompt
     type: string
@@ -17,9 +17,24 @@ parameters:
       zh_Hans: 提示词
     human_description:
       en_US: The text prompt used to generate the image.
-      zh_Hans: 用于生成图片的文字提示词
+      zh_Hans: 建议用英文的生成图片提示词以获得更好的生成效果。
     llm_description: this prompt text will be used to generate image.
     form: llm
+  - name: model
+    type: select
+    required: true
+    options:
+      - value: schnell
+        label:
+          en_US: Flux.1-schnell
+      - value: dev
+        label:
+          en_US: Flux.1-dev
+    default: schnell
+    label:
+      en_US: Choose Image Model
+      zh_Hans: 选择生成图片的模型
+    form: form
   - name: image_size
     type: select
     required: true