feat: add flux dev of siliconflow image-gen tool (#8450)

langgenius · Sep 15, 2024 · 3d083b7 · 3d083b7
1 parent aa5b2db
commit 3d083b7
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 5 deletions.
diff --git 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
@@ -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