From 9ce38e7db3379c08430052a9e812916809a89354 Mon Sep 17 00:00:00 2001
From: Will Miao <13051207myq@gmail.com>
Date: Thu, 20 Feb 2025 22:37:40 +0800
Subject: [PATCH] Add lora loader node

---
 __init__.py                            |  6 +--
 nodes/lora_gateway.py                  | 15 ------
 nodes/lora_loader.py                   | 75 ++++++++++++++++++++++++++
 static/js/components/ContextMenu.js    | 39 ++++++++++++++
 templates/components/context_menu.html |  5 ++
 5 files changed, 122 insertions(+), 18 deletions(-)
 delete mode 100644 nodes/lora_gateway.py
 create mode 100644 nodes/lora_loader.py

diff --git a/__init__.py b/__init__.py
index dfe3c0fe..d9083441 100644
--- a/__init__.py
+++ b/__init__.py
@@ -1,12 +1,12 @@
 from .lora_manager import LoraManager
-from .nodes.lora_gateway import LoRAGateway
+from .nodes.lora_loader import LoraManagerLoader
 
 NODE_CLASS_MAPPINGS = {
-    "LoRAGateway": LoRAGateway
+    "LoRALoader": LoraManagerLoader
 }
 
 NODE_DISPLAY_NAME_MAPPINGS = {
-    "LoRAGateway": "LoRAGateway"
+    "LoRALoader": "Lora Loader (LoraManager)"
 }
 
 WEB_DIRECTORY = "./web/comfyui"
diff --git a/nodes/lora_gateway.py b/nodes/lora_gateway.py
deleted file mode 100644
index 94a80d93..00000000
--- a/nodes/lora_gateway.py
+++ /dev/null
@@ -1,15 +0,0 @@
-class LoRAGateway:
-    """
-    LoRA Gateway Node
-    Acts as the entry point for LoRA management services
-    """
-    @classmethod
-    def INPUT_TYPES(cls):
-        return {
-            "required": {},
-            "optional": {}
-        }
-
-    RETURN_TYPES = ()
-    FUNCTION = "register_services"
-    CATEGORY = "LoRA Management"
\ No newline at end of file
diff --git a/nodes/lora_loader.py b/nodes/lora_loader.py
new file mode 100644
index 00000000..0c06ed0b
--- /dev/null
+++ b/nodes/lora_loader.py
@@ -0,0 +1,75 @@
+import re
+from nodes import LoraLoader
+from comfy.comfy_types import IO
+from ..services.lora_scanner import LoraScanner
+from ..config import config
+import asyncio
+import os
+
+class LoraManagerLoader:
+    NAME = "Lora Loader"
+    CATEGORY = "loaders"
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "model": ("MODEL",),
+                "clip": ("CLIP",),
+                "text": (IO.STRING, {
+                    "multiline": True, 
+                    "dynamicPrompts": True, 
+                    "tooltip": "Format: <lora:lora_name:strength> separated by spaces or punctuation"
+                }),
+            },
+        }
+
+    RETURN_TYPES = ("MODEL", "CLIP", IO.STRING, IO.STRING)
+    RETURN_NAMES = ("MODEL", "CLIP", "loaded_loras", "trigger_words")
+    FUNCTION = "load_loras"
+
+    async def get_lora_info(self, lora_name):
+        """Get the lora path and trigger words from cache"""
+        scanner = await LoraScanner.get_instance()
+        cache = await scanner.get_cached_data()
+        
+        for item in cache.raw_data:
+            if item.get('file_name') == lora_name:
+                file_path = item.get('file_path')
+                if file_path:
+                    for root in config.loras_roots:
+                        root = root.replace(os.sep, '/')
+                        if file_path.startswith(root):
+                            relative_path = os.path.relpath(file_path, root).replace(os.sep, '/')
+                            # Get trigger words from civitai metadata
+                            civitai = item.get('civitai', {})
+                            trigger_words = civitai.get('trainedWords', []) if civitai else []
+                            return relative_path, trigger_words
+        return lora_name, []  # Fallback if not found
+
+    def load_loras(self, model, clip, text):
+        """Loads multiple LoRAs based on the text input format."""
+        lora_pattern = r'<lora:([^:]+):([\d\.]+)>'
+        lora_matches = re.finditer(lora_pattern, text)
+        
+        loaded_loras = []
+        all_trigger_words = []
+        
+        for match in lora_matches:
+            lora_name = match.group(1)
+            strength = float(match.group(2))
+            
+            # Get lora path and trigger words
+            lora_path, trigger_words = asyncio.run(self.get_lora_info(lora_name))
+            
+            # Apply the LoRA using the resolved path
+            model, clip = LoraLoader().load_lora(model, clip, lora_path, strength, strength)
+            loaded_loras.append(f"{lora_name}: {strength}")
+            
+            # Add trigger words to collection
+            all_trigger_words.extend(trigger_words)
+        
+        loaded_loras_text = "\n".join(loaded_loras) if loaded_loras else "No LoRAs loaded"
+        trigger_words_text = ", ".join(all_trigger_words) if all_trigger_words else ""
+
+        return (model, clip, loaded_loras_text, trigger_words_text)
\ No newline at end of file
diff --git a/static/js/components/ContextMenu.js b/static/js/components/ContextMenu.js
index 7cb13884..4cd895d2 100644
--- a/static/js/components/ContextMenu.js
+++ b/static/js/components/ContextMenu.js
@@ -58,6 +58,9 @@ export class LoraContextMenu {
                 case 'refresh-metadata':
                     refreshSingleLoraMetadata(this.currentCard.dataset.filepath);
                     break;
+                case 'copy-to-stack':
+                    this.copyToLoraStack();
+                    break;
             }
             
             this.hideMenu();
@@ -98,4 +101,40 @@ export class LoraContextMenu {
         this.menu.style.display = 'none';
         this.currentCard = null;
     }
+
+    copyToLoraStack() {
+        if (!this.currentCard) return;
+
+        const loraStackNode = {
+            "id": crypto.randomUUID(),
+            "type": "LoRAStack",
+            "inputs": {
+                "enabled": true,
+                "lora_name": this.currentCard.dataset.filepath,
+                "model_strength": 1.0,  
+            },
+            "class_type": "LoRAStack",
+            "_meta": {
+                "title": `LoRA Stack (${this.currentCard.dataset.file_name})`,
+            }
+        };
+
+        // Convert to ComfyUI workflow format
+        const workflow = {
+            "last_node_id": 1,
+            "last_link_id": 0,
+            "nodes": [loraStackNode],
+            "links": [],
+        };
+
+        // Copy to clipboard
+        navigator.clipboard.writeText(JSON.stringify(workflow))
+            .then(() => {
+                showToast('LoRA Stack copied to clipboard', 'success');
+            })
+            .catch(err => {
+                console.error('Failed to copy:', err);
+                showToast('Failed to copy LoRA Stack', 'error');
+            });
+    }
 }
\ No newline at end of file
diff --git a/templates/components/context_menu.html b/templates/components/context_menu.html
index 8f4c1daa..4900171a 100644
--- a/templates/components/context_menu.html
+++ b/templates/components/context_menu.html
@@ -11,6 +11,11 @@
     <div class="context-menu-item" data-action="copyname">
         <i class="fas fa-copy"></i> Copy Model Name
     </div>
+    <div class="context-menu-separator"></div>
+    <div class="context-menu-item" data-action="copy-to-stack">
+        <i class="fas fa-layer-group"></i> Copy to LoRA Stack
+    </div>
+    <div class="context-menu-separator"></div>
     <div class="context-menu-item" data-action="preview">
         <i class="fas fa-image"></i> Replace Preview
     </div>