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..a54cd49 100644 --- a/nodes.py +++ b/nodes.py @@ -1,6 +1,9 @@ # (c) City96 || Apache-2.0 (apache.org/licenses/LICENSE-2.0) +import torch import gguf +import types import logging +import numpy as np import comfy.sd import comfy.utils @@ -20,7 +23,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 @@ -54,8 +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,) + 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: + return self.calculate_weight_orig(patches, weight, key) + + # 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) + + # 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.clone_orig = model.clone + model.clone = types.MethodType(clone, model) + + return model + NODE_CLASS_MAPPINGS = { "UnetLoaderGGUF": UnetLoaderGGUF, } diff --git a/ops.py b/ops.py index 83aa657..f52faf1 100644 --- a/ops.py +++ b/ops.py @@ -10,23 +10,28 @@ 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 + self.patches = patches - 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) 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 @@ -68,15 +73,48 @@ 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) 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) + + # apply patches + if patch_list: + weight = function(patch_list, weight, "dequant.name.unknown") + 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):