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):