From 0755dc14627cf37d9babda62d8fff8ea2c0b68c7 Mon Sep 17 00:00:00 2001
From: shadowcz007 <chizhiwei007@163.com>
Date: Wed, 24 Jan 2024 22:59:49 +0800
Subject: [PATCH] CreateLoraNames

---
 __init__.py                     |  8 +++++---
 data/extension-node-map.json    |  1 +
 nodes/Utils.py                  | 31 +++++++++++++++++++++++++++++++
 web/javascript/smart_connect.js | 20 ++++++++++++++++----
 4 files changed, 53 insertions(+), 7 deletions(-)

diff --git a/__init__.py b/__init__.py
index 824ddb9..2c48fed 100644
--- a/__init__.py
+++ b/__init__.py
@@ -557,7 +557,7 @@ from .nodes.ScreenShareNode import ScreenShareNode,FloatingVideo
 from .nodes.Clipseg import CLIPSeg,CombineMasks
 from .nodes.ChatGPT import ChatGPTNode,ShowTextForGPT,CharacterInText
 from .nodes.Audio import GamePal,SpeechRecognition,SpeechSynthesis
-from .nodes.Utils import CreateSampler_names,CreateCkptNames,CreateSeedNode,TESTNODE_,AppInfo,IntNumber,FloatSlider,TextInput,ColorInput,FontInput,TextToNumber,DynamicDelayProcessor,LimitNumber,SwitchByIndex,MultiplicationNode
+from .nodes.Utils import CreateLoraNames,CreateSampler_names,CreateCkptNames,CreateSeedNode,TESTNODE_,AppInfo,IntNumber,FloatSlider,TextInput,ColorInput,FontInput,TextToNumber,DynamicDelayProcessor,LimitNumber,SwitchByIndex,MultiplicationNode
 from .nodes.Mask import OutlineMask,FeatheredMask
 
 
@@ -620,7 +620,8 @@ NODE_CLASS_MAPPINGS = {
     "JoinWithDelimiter":JoinWithDelimiter,
     "Seed_":CreateSeedNode,
     "CkptNames_":CreateCkptNames,
-    "SamplerNames_":CreateSampler_names
+    "SamplerNames_":CreateSampler_names,
+    "LoraNames_":CreateLoraNames
     # "LaMaInpainting":LaMaInpainting
     # "GamePal":GamePal
 }
@@ -647,7 +648,8 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "PromptGenerate_Mix":"PromptGenerate ♾️Mixlab",
     "ChinesePrompt_Mix":"ChinesePrompt ♾️Mixlab",
     "GamePal":"GamePal ♾️Mixlab",
-    "RembgNode_Mix":"Removebg"
+    "RembgNode_Mix":"Removebg",
+    "LoraNames_":"LoraName_TriggerWords.safetensors"
 }
 
 # web ui的节点功能
diff --git a/data/extension-node-map.json b/data/extension-node-map.json
index d48903f..54ca2f1 100644
--- a/data/extension-node-map.json
+++ b/data/extension-node-map.json
@@ -4781,6 +4781,7 @@
             "Seed_",
             "CkptNames_",
             "SamplerNames_",
+            "LoraNames_",
             "CombineMasks_",
             "EnhanceImage",
             "GradientImage",
diff --git a/nodes/Utils.py b/nodes/Utils.py
index 8abc81e..5a14430 100644
--- a/nodes/Utils.py
+++ b/nodes/Utils.py
@@ -768,7 +768,38 @@ class CreateCkptNames:
         ckpt_names=ckpt_names.split('\n')
         ckpt_names = [name for name in ckpt_names if name.strip()]
         return (ckpt_names,)
+
+
+class CreateLoraNames:
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "lora_names": ("STRING",{"multiline": True,"default": "\n".join(folder_paths.get_filename_list("loras")),"dynamicPrompts": False}),
+            }
+        }
+
+    RETURN_TYPES = (any_type,"STRING",)
+    RETURN_NAMES = ("lora_names","prompt",)
     
+    INPUT_IS_LIST = False
+    OUTPUT_IS_LIST = (True,True,)
+
+    # OUTPUT_NODE = True
+    FUNCTION = "run"
+
+    CATEGORY = "♾️Mixlab/Utils"
+
+    def run(self, lora_names):
+        lora_names=lora_names.split('\n')
+        lora_names = [name for name in lora_names if name.strip()]
+        prompts=[os.path.splitext(n)[0] for n in lora_names]
+        return (lora_names,prompts,)
+   
+
 
 class CreateSampler_names:
     def __init__(self):
diff --git a/web/javascript/smart_connect.js b/web/javascript/smart_connect.js
index 56f7104..c1ff413 100644
--- a/web/javascript/smart_connect.js
+++ b/web/javascript/smart_connect.js
@@ -29,6 +29,18 @@ const smart_connect_config_input = [
     inputNodeName: 'SamplerNames_',
     inputNode_output_name: 'sampler_names'
   },
+  {
+    node_type: 'LoraLoaderModelOnly',
+    node_widget_name: 'lora_name',
+    inputNodeName: 'LoraNames_',
+    inputNode_output_name: 'lora_names'
+  },
+  {
+    node_type: 'LoadLoRA',
+    node_widget_name: 'lora_name',
+    inputNodeName: 'LoraNames_',
+    inputNode_output_name: 'lora_names'
+  },
   {
     node_type: 'Moondream',
     node_widget_name: 'image',
@@ -224,9 +236,9 @@ export function addSmartMenu (options, node) {
       // CLIPTextEncode 的widget ，name== 'text'
       let node_widget_name = sc.node_widget_name
       let widget = node.widgets.filter(w => w.name === node_widget_name)[0]
-      if(!widget){
+      if (!widget) {
         // 控件没有，则查找inputs
-        widget= node.inputs.filter(w => w.name === node_widget_name)[0]
+        widget = node.inputs.filter(w => w.name === node_widget_name)[0]
       }
 
       let isLinkNull = true
@@ -239,7 +251,7 @@ export function addSmartMenu (options, node) {
 
       if (widget && isLinkNull) {
         sopts.push({
-          content: sc.inputNodeName.split('_')[0]+'➡️',
+          content: sc.inputNodeName.split('_')[0] + '➡️',
           callback: () => {
             LGraphCanvas.prototype._createNodeForInput(
               node, //当前node
@@ -270,7 +282,7 @@ export function addSmartMenu (options, node) {
 
       if (widget && isLinkNull) {
         sopts.push({
-          content: '➡️'+ sc.outputNodeName.split('_')[0],
+          content: '➡️' + sc.outputNodeName.split('_')[0],
           callback: () => {
             LGraphCanvas.prototype._createNodeForOutput(
               node, //当前node