diff --git a/lora.py b/lora.py new file mode 100644 index 0000000..715fde8 --- /dev/null +++ b/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) diff --git a/nodes.py b/nodes.py index 8d63cae..7a720b5 100644 --- a/nodes.py +++ b/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, } diff --git a/ops.py b/ops.py index 95981b6..6e331de 100644 --- a/ops.py +++ b/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):