From 2877c185731a3801539f40861dcf29c6c4b7b243 Mon Sep 17 00:00:00 2001 From: City <125218114+city96@users.noreply.github.com> Date: Sat, 17 Aug 2024 04:42:59 +0200 Subject: [PATCH 1/4] Less fragile logic --- dequant.py | 3 --- nodes.py | 10 +++++++++- ops.py | 23 +++++++++++++---------- 3 files changed, 22 insertions(+), 14 deletions(-) diff --git a/dequant.py b/dequant.py index 8c38d30..8ab22ad 100644 --- a/dequant.py +++ b/dequant.py @@ -4,9 +4,6 @@ import torch import numpy as np def dequantize_tensor(tensor, dtype=torch.float16): - if tensor is None: - return None - data = torch.tensor(tensor.data) qtype = tensor.tensor_type oshape = tensor.tensor_shape diff --git a/nodes.py b/nodes.py index 3ce670c..8d63cae 100644 --- a/nodes.py +++ b/nodes.py @@ -1,6 +1,8 @@ # (c) City96 || Apache-2.0 (apache.org/licenses/LICENSE-2.0) +import torch import gguf import logging +import numpy as np import comfy.sd import comfy.utils @@ -20,7 +22,13 @@ def gguf_sd_loader(path): sd = {} dt = {} for tensor in reader.tensors: - sd[str(tensor.name)] = GGMLTensor(tensor) + sd[str(tensor.name)] = GGMLTensor( + torch.from_numpy(tensor.data), # mmap + tensor_type = tensor.tensor_type, + tensor_shape = torch.Size( + np.flip(list(tensor.shape)) + ) + ) dt[str(tensor.tensor_type)] = dt.get(str(tensor.tensor_type), 0) + 1 # sanity check debug print diff --git a/ops.py b/ops.py index 83aa657..95981b6 100644 --- a/ops.py +++ b/ops.py @@ -10,16 +10,13 @@ class GGMLTensor(torch.Tensor): """ Main tensor-like class for storing quantized weights """ - def __init__(self, tensor, *args, **kwargs): + def __init__(self, *args, tensor_type, tensor_shape, patches=[], **kwargs): super().__init__() - self.tensor_type = tensor.tensor_type - self.tensor_shape = torch.Size( - np.flip(list(tensor.shape)) - ) + self.tensor_type = tensor_type + self.tensor_shape = tensor_shape - def __new__(cls, tensor, *args, **kwargs): - data = torch.from_numpy(tensor.data) - return super().__new__(cls, data, *args, **kwargs) + def __new__(cls, *args, tensor_type, tensor_shape, patches=[], **kwargs): + return super().__new__(cls, *args, **kwargs) def to(self, *args, **kwargs): new = super().to(*args, **kwargs) @@ -74,9 +71,15 @@ class GGMLLayer(torch.nn.Module): super()._apply(fn) return self + def get_weight(self, tensor, dtype): + if tensor is None: + return + weight = dequantize_tensor(tensor, dtype) + return weight + def get_weights(self, dtype=torch.float16): - weight = dequantize_tensor(self.weight, dtype) - bias = dequantize_tensor(self.bias, dtype) + weight = self.get_weight(self.weight, dtype) + bias = self.get_weight(self.bias, dtype) return (weight, bias) class GGMLOps(comfy.ops.disable_weight_init): From fed8ff1e702ff5266890debe1dd472c568a11d37 Mon Sep 17 00:00:00 2001 From: City <125218114+city96@users.noreply.github.com> Date: Sat, 17 Aug 2024 05:28:46 +0200 Subject: [PATCH 2/4] Add LoRA support Initial, will need to change weight function --- lora.py | 37 +++++++++++++++++++++++++++++++++++++ nodes.py | 42 ++++++++++++++++++++++++++++++++++++++++++ ops.py | 16 +++++++++++++++- 3 files changed, 94 insertions(+), 1 deletion(-) create mode 100644 lora.py 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): From 22e2c885f00a5ba8a0acd439a322e16c6397911a Mon Sep 17 00:00:00 2001 From: City <125218114+city96@users.noreply.github.com> Date: Sat, 17 Aug 2024 07:38:55 +0200 Subject: [PATCH 3/4] Async patch load with GPU support --- nodes.py | 23 ++++++++++++++++++++++- ops.py | 25 +++++++++++++++++++++++-- 2 files changed, 45 insertions(+), 3 deletions(-) diff --git a/nodes.py b/nodes.py index 7a720b5..12b9212 100644 --- a/nodes.py +++ b/nodes.py @@ -76,6 +76,7 @@ class LoraLoaderGGUFModelOnly: "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}), + "lora_on_gpu": ("BOOLEAN", {"default": True}), } } @@ -84,7 +85,7 @@ class LoraLoaderGGUFModelOnly: CATEGORY = "bootleg" TITLE = "LoRA Loader Model Only (GGUF)" - def load_lora(self, model, lora_name, strength_model): + def load_lora(self, model, lora_name, strength_model, lora_on_gpu): if strength_model == 0: return (model,) @@ -102,6 +103,26 @@ class LoraLoaderGGUFModelOnly: lora = comfy.utils.load_torch_file(lora_path, safe_load=True) self.loaded_lora = (lora_path, lora) + # This is a super hacky method to get it to load alongside the model + param_dict_name = f"lora_state_dict"#_{lora_name}".replace('.', '_') + target_model = model.model.diffusion_model + if lora_on_gpu: + print(f"WARNING! Loading the LoRA to GPU memory may cause OOM issues or memory leaks!") + temp = {} + params = torch.nn.ParameterDict() + for k,v in lora.items(): + param = torch.nn.Parameter(v, requires_grad=False) + temp[k] = param + params[k.replace('.', '_')] = param + params.to(model.load_device, dtype=torch.float16) + target_model.register_module(param_dict_name, params) + lora = temp + else: + params = getattr(target_model, param_dict_name, None) + if params: + params.to("cpu", dtype=torch.float16) + delattr(target_model, param_dict_name) + model_lora = load_lora_gguf(model, lora, strength_model) return (model_lora,) diff --git a/ops.py b/ops.py index 6e331de..f52faf1 100644 --- a/ops.py +++ b/ops.py @@ -83,12 +83,33 @@ class GGMLLayer(torch.nn.Module): super()._apply(fn) return self + def move_patch_to_cuda(self, item, device): + if isinstance(item, torch.Tensor): + return item.to(device, non_blocking=True) + elif isinstance(item, tuple): + return tuple(self.move_patch_to_cuda(x, device) for x in item) + elif isinstance(item, list): + return [self.move_patch_to_cuda(x, device) for x in item] + else: + return item + def get_weight(self, tensor, dtype): if tensor is None: return + + # consolidate and load patches to GPU in async + patch_list = [] + device = tensor.device + t_move = lambda x: x.to(device) if torch.is_tensor(x) else x + for function, patches, _ in getattr(tensor, "patches", []): + patch_list += self.move_patch_to_cuda(patches, device) + + # dequantize tensor while patches load weight = dequantize_tensor(tensor, dtype) - for function, patches, key in getattr(tensor, "patches", []): - weight = function(patches, weight, key) + + # apply patches + if patch_list: + weight = function(patch_list, weight, "dequant.name.unknown") return weight def get_weights(self, dtype=torch.float16): From 96b0ba1021dfee589bdfeedc0c051cff85b613b6 Mon Sep 17 00:00:00 2001 From: City <125218114+city96@users.noreply.github.com> Date: Sat, 17 Aug 2024 07:56:38 +0200 Subject: [PATCH 4/4] Allow using the native LoRA loader node --- lora.py | 37 ------------------------- nodes.py | 84 ++++++++++++++++++-------------------------------------- 2 files changed, 26 insertions(+), 95 deletions(-) delete mode 100644 lora.py diff --git a/lora.py b/lora.py deleted file mode 100644 index 715fde8..0000000 --- a/lora.py +++ /dev/null @@ -1,37 +0,0 @@ -# (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 12b9212..a54cd49 100644 --- a/nodes.py +++ b/nodes.py @@ -1,6 +1,7 @@ # (c) City96 || Apache-2.0 (apache.org/licenses/LICENSE-2.0) import torch import gguf +import types import logging import numpy as np @@ -10,7 +11,6 @@ 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") @@ -63,70 +63,38 @@ class UnetLoaderGGUF: if model is None: logging.error("ERROR UNSUPPORTED UNET {}".format(unet_path)) raise RuntimeError("ERROR: Could not detect model type of: {}".format(unet_path)) + self.patch_calculate_weight(model) 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}), - "lora_on_gpu": ("BOOLEAN", {"default": True}), - } - } - - RETURN_TYPES = ("MODEL",) - FUNCTION = "load_lora" - CATEGORY = "bootleg" - TITLE = "LoRA Loader Model Only (GGUF)" - - def load_lora(self, model, lora_name, strength_model, lora_on_gpu): - 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] + def patch_calculate_weight(self, model): + 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: - temp = self.loaded_lora - self.loaded_lora = None - del temp + return self.calculate_weight_orig(patches, weight, key) - if lora is None: - lora = comfy.utils.load_torch_file(lora_path, safe_load=True) - self.loaded_lora = (lora_path, lora) + # make sure custom patch logic is used as to not mangle quants + model.calculate_weight_orig = model.calculate_weight + model.calculate_weight = types.MethodType(calculate_weight, model) - # This is a super hacky method to get it to load alongside the model - param_dict_name = f"lora_state_dict"#_{lora_name}".replace('.', '_') - target_model = model.model.diffusion_model - if lora_on_gpu: - print(f"WARNING! Loading the LoRA to GPU memory may cause OOM issues or memory leaks!") - temp = {} - params = torch.nn.ParameterDict() - for k,v in lora.items(): - param = torch.nn.Parameter(v, requires_grad=False) - temp[k] = param - params[k.replace('.', '_')] = param - params.to(model.load_device, dtype=torch.float16) - target_model.register_module(param_dict_name, params) - lora = temp - else: - params = getattr(target_model, param_dict_name, None) - if params: - params.to("cpu", dtype=torch.float16) - delattr(target_model, param_dict_name) + # make sure this logic is preserved for cloned patchers + def clone(self, *args, **kwargs): + new = self.clone_orig(*args, **kwargs) + new.calculate_weight_orig = new.calculate_weight + new.calculate_weight = types.MethodType(calculate_weight, model) + return new - model_lora = load_lora_gguf(model, lora, strength_model) - return (model_lora,) + model.clone_orig = model.clone + model.clone = types.MethodType(clone, model) + + return model NODE_CLASS_MAPPINGS = { "UnetLoaderGGUF": UnetLoaderGGUF, - "LoraLoaderGGUFModelOnly": LoraLoaderGGUFModelOnly, }