diff --git a/dequant.py b/dequant.py index f23aca6..949e5d6 100644 --- a/dequant.py +++ b/dequant.py @@ -7,15 +7,11 @@ from tqdm import tqdm def dequantize_tensor(tensor, dtype=torch.float16): data = torch.tensor(tensor.data) qtype = getattr(tensor, "tensor_type", None) - oshape = getattr(tensor, "tensor_shape", tensor.data.shape) + oshape = getattr(tensor, "tensor_shape", data.shape) - if qtype == None: - return data.to(dtype) - elif qtype == gguf.GGMLQuantizationType.F32: - return data.to(dtype) - elif qtype == gguf.GGMLQuantizationType.F16: - return data.to(dtype) - elif qtype in dequantize_functions: + if qtype in {None, gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16}: + return data.to(dtype).view(oshape) + if qtype in dequantize_functions: # dequantize in fp16 then convert instead of keeping FP32 out = dequantize(data, qtype, oshape, dtype=None) return out.to(dtype) if out.dtype != dtype else out # why is .to() not a no-op? diff --git a/nodes.py b/nodes.py index 628f66c..97686bd 100644 --- a/nodes.py +++ b/nodes.py @@ -28,13 +28,24 @@ def gguf_sd_loader(path): Read state dict as fake tensors """ reader = gguf.GGUFReader(path) + fixup_field = reader.get_field("comfy.gguf.tensor_fixup") + tensor_fixup = {} + if fixup_field is not None: + tensor_fixup_str=str(fixup_field.parts[fixup_field.data[-1]], encoding="utf-8") + for line in str(tensor_fixup_str).split("\n"): + key, *fixup_items = line.split(" ") + tensor_fixup[key] = torch.Size(tuple(int(i) for i in fixup_items)) sd = {} dt = {} for tensor in reader.tensors: - sd[str(tensor.name)] = GGMLTensor( + tensor_name = str(tensor.name) + shape = tensor_fixup.get(tensor_name) + if shape is None: + shape = torch.Size(tuple(int(v) for v in reversed(tensor.shape))) + sd[tensor_name] = GGMLTensor( torch.from_numpy(tensor.data), # mmap tensor_type = tensor.tensor_type, - tensor_shape = torch.Size(tuple(int(v) for v in reversed(tensor.shape))) + tensor_shape = shape ) dt[str(tensor.tensor_type)] = dt.get(str(tensor.tensor_type), 0) + 1 diff --git a/ops.py b/ops.py index 6bc1e2a..fac78d1 100644 --- a/ops.py +++ b/ops.py @@ -16,7 +16,7 @@ class GGMLTensor(torch.Tensor): self.tensor_shape = tensor_shape self.patches = patches - def __new__(cls, *args, tensor_type, tensor_shape, patches=[], **kwargs): + def __new__(cls, *args, tensor_type, tensor_shape, tensor_data=None, patches=[], **kwargs): return super().__new__(cls, *args, **kwargs) def to(self, *args, **kwargs): @@ -150,3 +150,26 @@ class GGMLOps(comfy.ops.manual_cast): if device: self.to(device) return x + + class Conv2d(GGMLLayer): + comfy_cast_weights = True + + def __init__(self, *args, device=None, dtype=None, **kwargs): + super().__init__(device=device, dtype=dtype) + _ = kwargs.pop("kernel_size", None) + self.kwargs=kwargs + + def forward(self, x): + # lowvram hack + device = None + if self.weight.device != x.device: + device = self.weight.device + self.to(x.device) + + weight, bias = self.get_weights(x.dtype) + x = torch.nn.functional.conv2d(x, weight, bias, **self.kwargs) + del weight, bias + + if device: + self.to(device) + return x diff --git a/tools/convert.py b/tools/convert.py index 178b664..df7f523 100644 --- a/tools/convert.py +++ b/tools/convert.py @@ -69,6 +69,9 @@ def handle_tensors(args, writer, state_dict): # - do something about this being awful and hacky max_name_len = max([len(s) for s in state_dict.keys()]) + 4 + if (max_name_len - 4) >= 64: + raise ValueError("Can only handle tensor names up to 63 characters!") + tensor_fixup = {} for key, data in tqdm(state_dict.items()): old_dtype = data.dtype @@ -122,6 +125,11 @@ def handle_tensors(args, writer, state_dict): #data = np.squeeze(data) # don't do this data_qtype = gguf.GGMLQuantizationType.F16 + if data_qtype != gguf.GGMLQuantizationType.F32 and n_dims >= 2 and n_params >= 256*256 and (n_params / 256).is_integer() and data.shape[0] >= 256 and data.shape[1] >= 256: + orig_shape=data.shape + orig_n_params=n_params + data = data.reshape(256+((n_params - 256*256) // 256),256) + tensor_fixup[key] = orig_shape try: data = gguf.quants.quantize(data, data_qtype) except gguf.QuantError as e: @@ -139,6 +147,8 @@ def handle_tensors(args, writer, state_dict): tqdm.write(f"{f'%-{max_name_len}s' % f'{new_name}'} {old_dtype} --> {data_qtype.name}, shape = {shape_str}") writer.add_tensor(new_name, data, raw_dtype=data_qtype) + fixup_str = "\n".join(" ".join((k,)+tuple(str(dimval) for dimval in v)) for k,v in tensor_fixup.items()) + writer.add_string("comfy.gguf.tensor_fixup", fixup_str) if __name__ == "__main__": args = parse_args()