Add LoRA support

Initial, will need to change weight function

lora.py

@@ -0,0 +1,37 @@
+# (c) City96 || Apache-2.0 (apache.org/licenses/LICENSE-2.0)
+import gguf
+import torch
+import types
+import logging
+
+import comfy.lora
+from .ops import GGMLTensor
+
+def load_lora_gguf(model, lora, strength):
+    key_map = comfy.lora.model_lora_keys_unet(model.model)
+    loaded = comfy.lora.load_lora(lora, key_map)
+
+    new = model.clone()
+    k = new.add_patches(loaded, strength)
+
+    if not getattr(new, "custom_calc_weight", False):
+        new.calculate_weight_orig = new.calculate_weight
+        new.calculate_weight = types.MethodType(calculate_weight, new)
+        new.custom_calc_weight = True
+
+    for key in [x for x in loaded if x not in k]:
+        logging.warning("NOT LOADED {}".format(key))
+
+    return (new)
+
+def calculate_weight(self, patches, weight, key):
+    if isinstance(weight, GGMLTensor):
+        qtype = weight.tensor_type
+        # TODO: don't even store these in a custom format
+        if qtype in [gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16]:
+            return self.calculate_weight_orig(patches, weight, key)
+        else:
+            weight.patches.append((self.calculate_weight_orig, patches, key))
+            return weight
+    else:
+        return self.calculate_weight_orig(patches, weight, key)

nodes.py

@@ -10,6 +10,7 @@ import comfy.model_management
 import folder_paths
 
 from .ops import GGMLTensor, GGMLOps
+from .lora import load_lora_gguf
 
 # TODO: This causes gguf files to show up in the main unet loader
 folder_paths.folder_names_and_paths["unet"][1].add(".gguf")
@@ -64,6 +65,47 @@ class UnetLoaderGGUF:
             raise RuntimeError("ERROR: Could not detect model type of: {}".format(unet_path))
         return (model,)
 
+class LoraLoaderGGUFModelOnly:
+    def __init__(self):
+        self.loaded_lora = None
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                 "model": ("MODEL",),
+                 "lora_name": (folder_paths.get_filename_list("loras"), ),
+                 "strength_model": ("FLOAT", {"default": 1.0, "min": -100.0, "max": 100.0, "step": 0.01}),
+            }
+        }
+
+    RETURN_TYPES = ("MODEL",)
+    FUNCTION = "load_lora"
+    CATEGORY = "bootleg"
+    TITLE = "LoRA Loader Model Only (GGUF)"
+
+    def load_lora(self, model, lora_name, strength_model):
+        if strength_model == 0:
+            return (model,)
+
+        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:
+            lora = comfy.utils.load_torch_file(lora_path, safe_load=True)
+            self.loaded_lora = (lora_path, lora)
+
+        model_lora = load_lora_gguf(model, lora, strength_model)
+        return (model_lora,)
+
 NODE_CLASS_MAPPINGS = {
     "UnetLoaderGGUF": UnetLoaderGGUF,
+    "LoraLoaderGGUFModelOnly": LoraLoaderGGUFModelOnly,
 }

ops.py

@@ -14,6 +14,7 @@ class GGMLTensor(torch.Tensor):
         super().__init__()
         self.tensor_type = tensor_type
         self.tensor_shape = tensor_shape
+        self.patches = patches
 
     def __new__(cls, *args, tensor_type, tensor_shape, patches=[], **kwargs):
         return super().__new__(cls, *args, **kwargs)
@@ -22,8 +23,15 @@ class GGMLTensor(torch.Tensor):
         new = super().to(*args, **kwargs)
         new.tensor_type = self.tensor_type
         new.tensor_shape = self.tensor_shape
+        new.patches = self.patches.copy()
         return new
 
+    def clone(self, *args, **kwargs):
+        return self
+
+    def detach(self, *args, **kwargs):
+        return self
+
     @property
     def shape(self):
         return self.tensor_shape
@@ -65,7 +73,11 @@ class GGMLLayer(torch.nn.Module):
 
     def _apply(self, fn):
         if self.weight is not None:
-            self.weight = fn(self.weight)
+            try:
+                self.weight = fn(self.weight)
+            except TypeError:
+                # TODO: Figure out why this happens
+                pass
         if self.bias is not None:
             self.bias = fn(self.bias)
         super()._apply(fn)
@@ -75,6 +87,8 @@ class GGMLLayer(torch.nn.Module):
         if tensor is None:
             return
         weight = dequantize_tensor(tensor, dtype)
+        for function, patches, key in getattr(tensor, "patches", []):
+            weight = function(patches, weight, key)
         return weight
 
     def get_weights(self, dtype=torch.float16):