From 24ba9fdf6c990ffa5d38463d38181c5415b418e0 Mon Sep 17 00:00:00 2001
From: zhuhao <37029601+hwzhuhao@users.noreply.github.com>
Date: Tue, 1 Oct 2024 10:55:54 +0800
Subject: [PATCH] feat: enhance stepfun image generation tool (#8954)

---
 .../provider/builtin/stepfun/tools/image.py   | 12 ++---
 .../provider/builtin/stepfun/tools/image.yaml | 50 ++++++++++---------
 2 files changed, 33 insertions(+), 29 deletions(-)

diff --git a/api/core/tools/provider/builtin/stepfun/tools/image.py b/api/core/tools/provider/builtin/stepfun/tools/image.py
index eb55dae51..61cc14fac 100644
--- a/api/core/tools/provider/builtin/stepfun/tools/image.py
+++ b/api/core/tools/provider/builtin/stepfun/tools/image.py
@@ -32,16 +32,17 @@ class StepfunTool(BuiltinTool):
         prompt = tool_parameters.get("prompt", "")
         if not prompt:
             return self.create_text_message("Please input prompt")
-
+        if len(prompt) > 1024:
+            return self.create_text_message("The prompt length should less than 1024")
         seed = tool_parameters.get("seed", 0)
         if seed > 0:
             extra_body["seed"] = seed
-        steps = tool_parameters.get("steps", 0)
+        steps = tool_parameters.get("steps", 50)
         if steps > 0:
             extra_body["steps"] = steps
-        negative_prompt = tool_parameters.get("negative_prompt", "")
-        if negative_prompt:
-            extra_body["negative_prompt"] = negative_prompt
+        cfg_scale = tool_parameters.get("cfg_scale", 7.5)
+        if cfg_scale > 0:
+            extra_body["cfg_scale"] = cfg_scale
 
         # call openapi stepfun model
         response = client.images.generate(
@@ -51,7 +52,6 @@ class StepfunTool(BuiltinTool):
             n=tool_parameters.get("n", 1),
             extra_body=extra_body,
         )
-        print(response)
 
         result = []
         for image in response.data:
diff --git a/api/core/tools/provider/builtin/stepfun/tools/image.yaml b/api/core/tools/provider/builtin/stepfun/tools/image.yaml
index 8d7c9b658..dfda6ed19 100644
--- a/api/core/tools/provider/builtin/stepfun/tools/image.yaml
+++ b/api/core/tools/provider/builtin/stepfun/tools/image.yaml
@@ -33,9 +33,9 @@ parameters:
     type: select
     required: false
     human_description:
-      en_US: used for selecting the image size
-      zh_Hans: 用于选择图像大小
-      pt_BR: used for selecting the image size
+      en_US: The size of the generated image
+      zh_Hans: 生成的图片大小
+      pt_BR: The size of the generated image
     label:
       en_US: Image size
       zh_Hans: 图像大小
@@ -77,17 +77,17 @@ parameters:
     type: number
     required: true
     human_description:
-      en_US: used for selecting the number of images
-      zh_Hans: 用于选择图像数量
-      pt_BR: used for selecting the number of images
+      en_US: Number of generated images, now only one image can be generated at a time
+      zh_Hans: 生成的图像数量，当前仅支持每次生成一张图片
+      pt_BR: Number of generated images, now only one image can be generated at a time
     label:
-      en_US: Number of images
-      zh_Hans: 图像数量
-      pt_BR: Number of images
+      en_US: Number of generated images
+      zh_Hans: 生成的图像数量
+      pt_BR: Number of generated images
     form: form
     default: 1
     min: 1
-    max: 10
+    max: 1
   - name: seed
     type: number
     required: false
@@ -109,21 +109,25 @@ parameters:
       zh_Hans: Steps
       pt_BR: Steps
     human_description:
-      en_US: Steps
-      zh_Hans: Steps
-      pt_BR: Steps
+      en_US: Steps, now support integers between 1 and 100
+      zh_Hans: Steps, 当前支持 1～100 之间整数
+      pt_BR: Steps, now support integers between 1 and 100
     form: form
-    default: 10
-  - name: negative_prompt
-    type: string
+    default: 50
+    min: 1
+    max: 100
+  - name: cfg_scale
+    type: number
     required: false
     label:
-      en_US: Negative prompt
-      zh_Hans: Negative prompt
-      pt_BR: Negative prompt
+      en_US: classifier-free guidance scale
+      zh_Hans: classifier-free guidance scale
+      pt_BR: classifier-free guidance scale
     human_description:
-      en_US: Negative prompt
-      zh_Hans: Negative prompt
-      pt_BR: Negative prompt
+      en_US: classifier-free guidance scale
+      zh_Hans: classifier-free guidance scale
+      pt_BR: classifier-free guidance scale
     form: form
-    default: (worst quality:1.3), (nsfw), low quality
+    default: 7.5
+    min: 1
+    max: 10