From ef308fd121ab8a52e3e96da7699000dc9d85db97 Mon Sep 17 00:00:00 2001
From: Yeuoly <45712896+Yeuoly@users.noreply.github.com>
Date: Thu, 1 Feb 2024 13:12:57 +0800
Subject: [PATCH] feat: add sd model parameter (#2331)

---
 .../stablediffusion/tools/stable_diffusion.py        |  4 ++++
 .../stablediffusion/tools/stable_diffusion.yaml      | 12 ++++++++++++
 2 files changed, 16 insertions(+)

diff --git a/api/core/tools/provider/builtin/stablediffusion/tools/stable_diffusion.py b/api/core/tools/provider/builtin/stablediffusion/tools/stable_diffusion.py
index dbf2fd749f..b9e91698ec 100644
--- a/api/core/tools/provider/builtin/stablediffusion/tools/stable_diffusion.py
+++ b/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')
diff --git a/api/core/tools/provider/builtin/stablediffusion/tools/stable_diffusion.yaml b/api/core/tools/provider/builtin/stablediffusion/tools/stable_diffusion.yaml
index 5404b84360..bd17f5098e 100644
--- a/api/core/tools/provider/builtin/stablediffusion/tools/stable_diffusion.yaml
+++ b/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