From 44d368ff17f013783b4c988f33c6bf44330c8b29 Mon Sep 17 00:00:00 2001
From: Martin Bukowski <mbukowski@bmds.us>
Date: Fri, 26 Jan 2024 19:03:26 -0600
Subject: [PATCH] lora tags

---
 README.md          |  5 +++++
 __init__.py        |  4 ++++
 components/lora.py | 54 ++++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 63 insertions(+)
 create mode 100644 components/lora.py

diff --git a/README.md b/README.md
index fb9fe90..990046c 100644
--- a/README.md
+++ b/README.md
@@ -24,6 +24,11 @@ Merge two checkpoint models by dare ties (https://github.com/yule-BUAA/MergeLM).
 |mask|Mask Operations|`MODEL_MASK`, `MODEL_MASK`|`MODEL_MASK`|Allows set operations to be performed on masks|
 |mask|Mask Edit|`MODEL_MASK`|`MODEL_MASK`|Allows the direct editing of mask layers|
 
+## LoRA
+|category|node name|input type|output type|desc.|
+| --- | --- | --- | --- | --- |
+|lora|LoRA Loader (Tags)|`MODEL`, `CLIP`|`MODEL`, `CLIP`, `STRING`|Loads a LoRA model, returning the tags from the metadata|
+
 ## Utilities
 |category|node name|input type|output type|desc.|
 | --- | --- | --- | --- | --- |
diff --git a/__init__.py b/__init__.py
index e9fd227..248b2a1 100644
--- a/__init__.py
+++ b/__init__.py
@@ -5,6 +5,7 @@ from .components.block import BlockUnetMerger
 from .components.normalize import NormalizeUnet
 from .components.mask_model import MagnitudeMasker, MaskOperations, MaskEdit, SimpleMasker, QuadMasker
 from .components.reports import MaskReporting, ModelReporting, LoRAReporting
+from .components.lora import LoraLoaderTags
 
 
 NODE_CLASS_MAPPINGS = {
@@ -21,6 +22,8 @@ NODE_CLASS_MAPPINGS = {
     "DM_MaskReporting": MaskReporting,
     "DM_LoRAReporting": LoRAReporting,
     "DM_NormalizeModel": NormalizeUnet,
+    "DM_LoRALoaderTags": LoraLoaderTags,
+    
 }
 
 NODE_DISPLAY_NAME_MAPPINGS = {
@@ -37,6 +40,7 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "DM_MaskReporting": "Mask Reporting",
     "DM_LoRAReporting": "LoRA Reporting",
     "DM_NormalizeModel": "Normalize Model",
+    "DM_LoRALoaderTags": "LoRA Loader (Tags)",
 }
 
 __all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']
diff --git a/components/lora.py b/components/lora.py
new file mode 100644
index 0000000..b0ac9e0
--- /dev/null
+++ b/components/lora.py
@@ -0,0 +1,54 @@
+# components/lora.py
+import folder_paths
+from comfy.sd import load_lora_for_models
+from comfy.sd import CLIP
+from comfy.model_patcher import ModelPatcher
+
+from ..ddare.lora import DoctorLora
+from ..ddare.const import LORA_CATEGORY
+
+class LoraLoaderTags:
+    def __init__(self):
+        self.loaded_lora = None
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required": { "model": ("MODEL",),
+                              "clip": ("CLIP", ),
+                              "lora_name": (folder_paths.get_filename_list("loras"), ),
+                              "strength_model": ("FLOAT", {"default": 1.0, "min": -20.0, "max": 20.0, "step": 0.01}),
+                              "strength_clip": ("FLOAT", {"default": 1.0, "min": -20.0, "max": 20.0, "step": 0.01}),
+                              "tag_limit": ("INT", {"default": 10, "min": 1, "max": 100}),
+                              }}
+    RETURN_TYPES = ("MODEL", "CLIP", "STRING")
+    FUNCTION = "load_lora"
+
+    CATEGORY = LORA_CATEGORY
+
+    def load_lora(self, model : ModelPatcher, clip : CLIP, lora_name : str, strength_model : float, strength_clip : float, tag_limit : int = 10, **kwargs):
+        tags = ""
+        if strength_model == 0 and strength_clip == 0:
+            return (model, clip, tags)
+
+        lora_path = folder_paths.get_full_path("loras", lora_name)
+        lora = None
+        if self.loaded_lora is not None:
+            if self.loaded_lora[0] == lora_path:
+                lora = self.loaded_lora[1]
+            else:
+                temp = self.loaded_lora
+                self.loaded_lora = None
+                del temp
+
+        if lora is None:
+            drlora = DoctorLora.load(lora_path)
+
+            tags = drlora.tags
+            if tags is not None:
+                tags = ", ".join(tags[:tag_limit])
+
+            lora = drlora.lora
+
+        model_lora, clip_lora = load_lora_for_models(model, clip, lora, strength_model, strength_clip)
+        return (model_lora, clip_lora, tags)
+ 
\ No newline at end of file