Explorar el Código

feat: add sd model parameter (#2331)

Yeuoly hace 1 año
padre
commit
ef308fd121

+ 4 - 0
api/core/tools/provider/builtin/stablediffusion/tools/stable_diffusion.py

@@ -70,6 +70,10 @@ class StableDiffusionTool(BuiltinTool):
         base_url = self.runtime.credentials.get('base_url', None)
         if not base_url:
             return self.create_text_message('Please input base_url')
+        
+        if 'model' in tool_parameters:
+            self.runtime.credentials['model'] = tool_parameters['model']
+        
         model = self.runtime.credentials.get('model', None)
         if not model:
             return self.create_text_message('Please input model')

+ 12 - 0
api/core/tools/provider/builtin/stablediffusion/tools/stable_diffusion.yaml

@@ -25,6 +25,18 @@ parameters:
       pt_BR: Image prompt, you can check the official documentation of Stable Diffusion
     llm_description: Image prompt of Stable Diffusion, you should describe the image you want to generate as a list of words as possible as detailed, the prompt must be written in English.
     form: llm
+  - name: model
+    type: string
+    required: false
+    label:
+      en_US: Model Name
+      zh_Hans: 模型名称
+      pt_BR: Model Name
+    human_description:
+      en_US: Model Name
+      zh_Hans: 模型名称
+      pt_BR: Model Name
+    form: form
   - name: lora
     type: string
     required: false