Merge pull request #29 from city96/lora

Native LoRA support

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

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,
 }

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