SD1/SDXL fixes

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?

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
 

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

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