From 45dd31814053b91c88beeed41956020cd88d3f76 Mon Sep 17 00:00:00 2001
From: TSC <112517630+LucianoCirino@users.noreply.github.com>
Date: Sat, 8 Jul 2023 10:29:21 -0500
Subject: [PATCH] Lora Loader Fix + Efficiency Loader Clip re-add

---
 efficiency_nodes.py | 14 +++++---------
 lora_patch.py       |  6 +++---
 2 files changed, 8 insertions(+), 12 deletions(-)

diff --git a/efficiency_nodes.py b/efficiency_nodes.py
index e6633e0..3e3ca1f 100644
--- a/efficiency_nodes.py
+++ b/efficiency_nodes.py
@@ -40,11 +40,7 @@ import comfy.sd
 import comfy.utils
 
 # Load legacy lora functions
-from lora_patch import load_lora, load_lora_for_models
-
-# Replace the lora functions with the legacy functions
-comfy.sd.load_lora = load_lora
-comfy.sd.load_lora_for_models = load_lora_for_models
+from lora_patch import load_lora_legacy, load_lora_for_models_legacy
 
 MAX_RESOLUTION=8192
 
@@ -300,7 +296,7 @@ def load_lora(lora_name, ckpt_name, strength_model, strength_clip, id, cache=Non
 
     ckpt, clip, _ = load_checkpoint(ckpt_name, id, cache=ckpt_cache, cache_overwrite=cache_overwrite)
     lora_path = folder_paths.get_full_path("loras", lora_name)
-    lora_model, lora_clip = comfy.sd.load_lora_for_models(ckpt, clip, lora_path, strength_model, strength_clip)
+    lora_model, lora_clip = load_lora_for_models_legacy(ckpt, clip, lora_path, strength_model, strength_clip)
 
     if cache:
         if len([entry for entry in loaded_objects["lora"] if id in entry[-1]]) < cache:
@@ -425,8 +421,8 @@ class TSC_EfficientLoader:
                             "my_unique_id": "UNIQUE_ID",},
                 }
 
-    RETURN_TYPES = ("MODEL", "CONDITIONING", "CONDITIONING", "LATENT", "VAE", "DEPENDENCIES",)
-    RETURN_NAMES = ("MODEL", "CONDITIONING+", "CONDITIONING-", "LATENT", "VAE", "DEPENDENCIES", )
+    RETURN_TYPES = ("MODEL", "CONDITIONING", "CONDITIONING", "LATENT", "VAE", "CLIP", "DEPENDENCIES",)
+    RETURN_NAMES = ("MODEL", "CONDITIONING+", "CONDITIONING-", "LATENT", "VAE", "CLIP", "DEPENDENCIES", )
     FUNCTION = "efficientloader"
     CATEGORY = "Efficiency Nodes/Loaders"
 
@@ -469,7 +465,7 @@ class TSC_EfficientLoader:
         # Data for XY Plot
         dependencies = (vae_name, ckpt_name, clip, clip_skip, positive, negative, lora_name, lora_model_strength, lora_clip_strength)
 
-        return (model, [[clip.encode(positive), {}]], [[clip.encode(negative), {}]], {"samples":latent}, vae, dependencies, )
+        return (model, [[clip.encode(positive), {}]], [[clip.encode(negative), {}]], {"samples":latent}, vae, clip, dependencies, )
 
 ########################################################################################################################
 # TSC KSampler (Efficient)
diff --git a/lora_patch.py b/lora_patch.py
index 5c57a12..205feac 100644
--- a/lora_patch.py
+++ b/lora_patch.py
@@ -46,7 +46,7 @@ LORA_UNET_MAP_RESNET = {
     "skip_connection": "resnets_{}_conv_shortcut"
 }
 
-def load_lora(path, to_load):
+def load_lora_legacy(path, to_load):
     lora = utils.load_torch_file(path)
     patch_dict = {}
     loaded_keys = set()
@@ -244,10 +244,10 @@ def model_lora_keys(model, key_map={}):
 
     return key_map
 
-def load_lora_for_models(model, clip, lora_path, strength_model, strength_clip):
+def load_lora_for_models_legacy(model, clip, lora_path, strength_model, strength_clip):
     key_map = model_lora_keys(model.model)
     key_map = model_lora_keys(clip.cond_stage_model, key_map)
-    loaded = load_lora(lora_path, key_map)
+    loaded = load_lora_legacy(lora_path, key_map)
     new_modelpatcher = model.clone()
     k = new_modelpatcher.add_patches(loaded, strength_model)
     new_clip = clip.clone()