Merge pull request #92 from city96/ops_refactor

Fix LoRA reloading, ops refactor, better lowvram speed, sd.cpp compat

dequant.py

@@ -3,21 +3,29 @@ import gguf
 import torch
 from tqdm import tqdm
 
-def dequantize_tensor(tensor, dtype=None, dequant_dtype=None):
-    data = torch.tensor(tensor.data)
-    qtype = getattr(tensor, "tensor_type", None)
-    oshape = getattr(tensor, "tensor_shape", data.shape)
 
-    if qtype in (None, gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16):
-        return data.to(dtype)
+TORCH_COMPATIBLE_QTYPES = {None, gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16}
+
+def is_torch_compatible(tensor):
+    return tensor is None or getattr(tensor, "tensor_type", None) in TORCH_COMPATIBLE_QTYPES
+
+def is_quantized(tensor):
+    return not is_torch_compatible(tensor)
+
+def dequantize_tensor(tensor, dtype=None, dequant_dtype=None):
+    qtype = getattr(tensor, "tensor_type", None)
+    oshape = getattr(tensor, "tensor_shape", tensor.shape)
+
+    if qtype in TORCH_COMPATIBLE_QTYPES:
+        return tensor.to(dtype)
     elif qtype in dequantize_functions:
         dequant_dtype = dtype if dequant_dtype == "target" else dequant_dtype
-        return dequantize(data, qtype, oshape, dtype=dequant_dtype).to(dtype)
+        return dequantize(tensor.data, qtype, oshape, dtype=dequant_dtype).to(dtype)
     else:
         # this is incredibly slow
         tqdm.write(f"Falling back to numpy dequant for qtype: {qtype}")
-        new = gguf.quants.dequantize(data.cpu().numpy(), qtype)
-        return torch.from_numpy(new).to(data.device, dtype=dtype)
+        new = gguf.quants.dequantize(tensor.cpu().numpy(), qtype)
+        return torch.from_numpy(new).to(tensor.device, dtype=dtype)
 
 def dequantize(data, qtype, oshape, dtype=None):
     """
@@ -45,6 +53,10 @@ def split_block_dims(blocks, *args):
     dims = list(args) + [n_max - sum(args)]
     return torch.split(blocks, dims, dim=1)
 
+# Full weights #
+def dequantize_blocks_BF16(blocks, block_size, type_size, dtype=None):
+    return (blocks.view(torch.int16).to(torch.int32) << 16).view(torch.float32)
+
 # Legacy Quants #
 def dequantize_blocks_Q8_0(blocks, block_size, type_size, dtype=None):
     d, x = split_block_dims(blocks, 2)
@@ -222,6 +234,7 @@ def dequantize_blocks_Q2_K(blocks, block_size, type_size, dtype=None):
     return qs.reshape((n_blocks, -1))
 
 dequantize_functions = {
+    gguf.GGMLQuantizationType.BF16: dequantize_blocks_BF16,
     gguf.GGMLQuantizationType.Q8_0: dequantize_blocks_Q8_0,
     gguf.GGMLQuantizationType.Q5_1: dequantize_blocks_Q5_1,
     gguf.GGMLQuantizationType.Q5_0: dequantize_blocks_Q5_0,

nodes.py

@@ -10,8 +10,8 @@ import comfy.model_management
 import comfy.model_patcher
 import folder_paths
 
-from .ops import GGMLTensor, GGMLOps, move_patch_to_cuda
-from .dequant import dequantize_tensor
+from .ops import GGMLTensor, GGMLOps, move_patch_to_device
+from .dequant import is_quantized, is_torch_compatible
 
 # Add a custom keys for files ending in .gguf
 if "unet_gguf" not in folder_paths.folder_names_and_paths:
@@ -32,34 +32,73 @@ def gguf_sd_loader_get_orig_shape(reader, tensor_name):
         raise TypeError(f"Bad original shape metadata for {field_key}: Expected ARRAY of INT32, got {field.types}")
     return torch.Size(tuple(int(field.parts[part_idx][0]) for part_idx in field.data))
 
-def gguf_sd_loader(path):
+def gguf_sd_loader(path, handle_prefix="model.diffusion_model."):
     """
     Read state dict as fake tensors
     """
     reader = gguf.GGUFReader(path)
-    sd = {}
-    dt = {}
+
+    # filter and strip prefix
+    has_prefix = False
+    if handle_prefix is not None:
+        prefix_len = len(handle_prefix)
+        tensor_names = set(tensor.name for tensor in reader.tensors)
+        has_prefix = any(s.startswith(handle_prefix) for s in tensor_names)
+
+    tensors = []
     for tensor in reader.tensors:
-        tensor_name = str(tensor.name)
+        sd_key = tensor_name = tensor.name
+        if has_prefix:
+            if not tensor_name.startswith(handle_prefix):
+                continue
+            sd_key = tensor_name[prefix_len:]
+        tensors.append((sd_key, tensor))
+
+    # detect and verify architecture
+    compat = None
+    arch_str = None
+    arch_field = reader.get_field("general.architecture")
+    if arch_field is not None:
+        if len(arch_field.types) != 1 or arch_field.types[0] != gguf.GGUFValueType.STRING:
+            raise TypeError(f"Bad type for GGUF general.architecture key: expected string, got {arch_field.types!r}")
+        arch_str = str(arch_field.parts[arch_field.data[-1]], encoding="utf-8")
+        if arch_str not in {"flux", "sd1", "sdxl", "t5", "t5encoder"}:
+            raise ValueError(f"Unexpected architecture type in GGUF file, expected one of flux, sd1, sdxl, t5encoder but got {arch_str!r}")
+    else: # stable-diffusion.cpp
+        # import here to avoid changes to convert.py breaking regular models
+        from .tools.convert import detect_arch
+        arch_str = detect_arch(set(val[0] for val in tensors))
+        compat = "sd.cpp"
+
+    # main loading loop
+    state_dict = {}
+    qtype_dict = {}
+    for sd_key, tensor in tensors:
+        tensor_name = tensor.name
+        tensor_type_str = str(tensor.tensor_type)
         torch_tensor = torch.from_numpy(tensor.data) # mmap
+
         shape = gguf_sd_loader_get_orig_shape(reader, tensor_name)
         if shape is None:
             shape = torch.Size(tuple(int(v) for v in reversed(tensor.shape)))
-        elif tensor.tensor_type in {gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16}:
+            # Workaround for stable-diffusion.cpp SDXL detection.
+            if compat == "sd.cpp" and arch_str == "sdxl":
+                if any([tensor_name.endswith(x) for x in (".proj_in.weight", ".proj_out.weight")]):
+                    while len(shape) > 2 and shape[-1] == 1:
+                        shape = shape[:-1]
+
+        # add to state dict
+        if tensor.tensor_type in {gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16}:
             torch_tensor = torch_tensor.view(*shape)
-        sd[tensor_name] = GGMLTensor(
-            torch_tensor,
-            tensor_type = tensor.tensor_type,
-            tensor_shape = shape
-        )
-        dt[str(tensor.tensor_type)] = dt.get(str(tensor.tensor_type), 0) + 1
+        state_dict[sd_key] = GGMLTensor(torch_tensor, tensor_type=tensor.tensor_type, tensor_shape=shape)
+        qtype_dict[tensor_type_str] = qtype_dict.get(tensor_type_str, 0) + 1
 
     # sanity check debug print
     print("\nggml_sd_loader:")
-    for k,v in dt.items():
+    for k,v in qtype_dict.items():
         print(f" {k:30}{v:3}")
-    print("\n")
-    return sd
+
+    return state_dict
 
 # for remapping llama.cpp -> original key names
 clip_sd_map = {
@@ -78,10 +117,6 @@ clip_sd_map = {
     "ffn_gate": "layer.1.DenseReluDense.wi_0",
     "ffn_norm": "layer.1.layer_norm",
 }
-# weights that should be dequantized on load
-clip_sd_dequant = {
-    "shared.weight",
-}
 
 def gguf_clip_loader(path):
     raw_sd = gguf_sd_loader(path)
@@ -90,9 +125,6 @@ def gguf_clip_loader(path):
     for k,v in raw_sd.items():
         for s,d in clip_sd_map.items():
             k = k.replace(s,d)
-        if k in clip_sd_dequant:
-            v = dequantize_tensor(v, torch.float32).to(torch.float16)
-            v = GGMLTensor(v, tensor_type=gguf.GGMLQuantizationType.F16, tensor_shape=v.shape)
         sd[k] = v
     return sd
 
@@ -105,9 +137,6 @@ class GGUFModelPatcher(comfy.model_patcher.ModelPatcher):
         if key not in self.patches:
             return
         weight = comfy.utils.get_attr(self.model, key)
-        inplace_update = self.weight_inplace_update or inplace_update
-        if key not in self.backup:
-            self.backup[key] = collections.namedtuple('Dimension', ['weight', 'inplace_update'])(weight.to(device=self.offload_device, copy=inplace_update), inplace_update)
 
         try:
             from comfy.lora import calculate_weight
@@ -115,18 +144,17 @@ class GGUFModelPatcher(comfy.model_patcher.ModelPatcher):
             calculate_weight = self.calculate_weight
 
         patches = self.patches[key]
-        qtype = getattr(weight, "tensor_type", None)
-        if qtype not in (None, gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16):
-            if device_to is not None:
-                out_weight = weight.to(device_to, copy=True)
-            else:
-                out_weight = weight.to(weight.device, copy=True) # make sure patches are copied
-
-            if self.patch_on_device:
-                patches = move_patch_to_cuda(patches, self.load_device)
+        if is_quantized(weight):
+            out_weight = weight.to(device_to)
+            patches = move_patch_to_device(patches, self.load_device if self.patch_on_device else self.offload_device)
             out_weight.patches.append((calculate_weight, patches, key))
-
         else:
+            inplace_update = self.weight_inplace_update or inplace_update
+            if key not in self.backup:
+                self.backup[key] = collections.namedtuple('Dimension', ['weight', 'inplace_update'])(
+                    weight.to(device=self.offload_device, copy=inplace_update), inplace_update
+                )
+
             if device_to is not None:
                 temp_weight = comfy.model_management.cast_to_device(weight, device_to, torch.float32, copy=True)
             else:
@@ -140,6 +168,16 @@ class GGUFModelPatcher(comfy.model_patcher.ModelPatcher):
         else:
             comfy.utils.set_attr_param(self.model, key, out_weight)
 
+    def unpatch_model(self, device_to=None, unpatch_weights=True):
+        if unpatch_weights:
+            for p in self.model.parameters():
+                if is_torch_compatible(p):
+                    continue
+                patches = getattr(p, "patches", [])
+                if len(patches) > 0:
+                    p.patches = []
+        return super().unpatch_model(device_to=None, unpatch_weights=unpatch_weights)
+
     def load(self, *args, force_patch_weights=False, **kwargs):
         # always call `patch_weight_to_device` even for lowvram
         return super().load(*args, force_patch_weights=True, **kwargs)
@@ -176,14 +214,14 @@ class UnetLoaderGGUF:
     def load_unet(self, unet_name, dequant_dtype=None, patch_dtype=None, patch_on_device=None):
         ops = GGMLOps()
 
-        if dequant_dtype in ["default", None]:
+        if dequant_dtype in ("default", None):
             ops.Linear.dequant_dtype = None
         elif dequant_dtype in ["target"]:
             ops.Linear.dequant_dtype = dequant_dtype
         else:
             ops.Linear.dequant_dtype = getattr(torch, dequant_dtype)
 
-        if patch_dtype in ["default", None]:
+        if patch_dtype in ("default", None):
             ops.Linear.patch_dtype = None
         elif patch_dtype in ["target"]:
             ops.Linear.patch_dtype = patch_dtype
@@ -287,11 +325,12 @@ class CLIPLoaderGGUF:
 class DualCLIPLoaderGGUF(CLIPLoaderGGUF):
     @classmethod
     def INPUT_TYPES(s):
+        file_options = (s.get_filename_list(), )
         return {
             "required": {
-                "clip_name1": (s.get_filename_list(), ),
-                "clip_name2": (s.get_filename_list(), ),
-                "type": (["sdxl", "sd3", "flux"], ),
+                "clip_name1": file_options,
+                "clip_name2": file_options,
+                "type": (("sdxl", "sd3", "flux"), ),
             }
         }
 
@@ -300,18 +339,19 @@ class DualCLIPLoaderGGUF(CLIPLoaderGGUF):
     def load_clip(self, clip_name1, clip_name2, type):
         clip_path1 = folder_paths.get_full_path("clip", clip_name1)
         clip_path2 = folder_paths.get_full_path("clip", clip_name2)
-        clip_paths = [clip_path1, clip_path2]
+        clip_paths = (clip_path1, clip_path2)
         clip_type = clip_name_dict.get(type, comfy.sd.CLIPType.STABLE_DIFFUSION)
         return (self.load_patcher(clip_paths, clip_type, self.load_data(clip_paths)),)
 
 class TripleCLIPLoaderGGUF(CLIPLoaderGGUF):
     @classmethod
     def INPUT_TYPES(s):
+        file_options = (s.get_filename_list(), )
         return {
             "required": {
-                "clip_name1": (s.get_filename_list(), ),
-                "clip_name2": (s.get_filename_list(), ),
-                "clip_name3": (s.get_filename_list(), ),
+                "clip_name1": file_options,
+                "clip_name2": file_options,
+                "clip_name3": file_options,
             }
         }
 
@@ -321,7 +361,7 @@ class TripleCLIPLoaderGGUF(CLIPLoaderGGUF):
         clip_path1 = folder_paths.get_full_path("clip", clip_name1)
         clip_path2 = folder_paths.get_full_path("clip", clip_name2)
         clip_path3 = folder_paths.get_full_path("clip", clip_name3)
-        clip_paths = [clip_path1, clip_path2, clip_path3]
+        clip_paths = (clip_path1, clip_path2, clip_path3)
         clip_type = clip_name_dict.get(type, comfy.sd.CLIPType.STABLE_DIFFUSION)
         return (self.load_patcher(clip_paths, clip_type, self.load_data(clip_paths)),)
 

ops.py

@@ -1,9 +1,10 @@
 # (c) City96 || Apache-2.0 (apache.org/licenses/LICENSE-2.0)
+import gguf
 import torch
-from functools import partial
 
 import comfy.ops
-from .dequant import dequantize_tensor
+import comfy.model_management
+from .dequant import dequantize_tensor, is_quantized
 
 class GGMLTensor(torch.Tensor):
     """
@@ -59,67 +60,49 @@ class GGMLLayer(torch.nn.Module):
     comfy_cast_weights = True
     dequant_dtype = None
     patch_dtype = None
+    torch_compatible_tensor_types = {None, gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16}
 
-    def __init__(self, *args, **kwargs):
-        super().__init__()
-        self.weight = GGMLTensor(1, tensor_type=None, tensor_shape=None)
-        self.bias = None
+    def is_ggml_quantized(self, *, weight=None, bias=None):
+        if weight is None:
+            weight = self.weight
+        if bias is None:
+            bias = self.bias
+        return is_quantized(weight) or is_quantized(bias)
 
-    def _forward_operation(self, x, weight, bias):
-        raise NotImplementedError
+    def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs):
+        weight, bias = state_dict.get(f"{prefix}weight"), state_dict.get(f"{prefix}bias")
+        if self.is_ggml_quantized(weight=weight, bias=bias):
+            return self.ggml_load_from_state_dict(state_dict, prefix, *args, **kwargs)
+        return super()._load_from_state_dict(state_dict, prefix, *args, **kwargs)
 
-    def forward(self, x):
-        # lowvram hack
-        device = None
-        if self.weight.device != x.device:
-            device = self.weight.device
-            self.to(x.device)
+    def _save_to_state_dict(self, *args, **kwargs):
+        if self.is_ggml_quantized():
+            return self.ggml_save_to_state_dict(*args, **kwargs)
+        return super()._save_to_state_dict(*args, **kwargs)
 
-        weight, bias = self.get_weights(x.dtype)
-        x = self._forward_operation(x, weight, bias=bias)
-        del weight, bias
-
-        if device:
-            self.to(device)
-        return x
-
-    def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs):
+    def ggml_load_from_state_dict(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs):
+        prefix_len = len(prefix)
         for k,v in state_dict.items():
-            if k[len(prefix):] == "weight":
-                self.weight = v
-            elif k[len(prefix):] == "bias":
-                self.bias = v
+            if k[prefix_len:] == "weight":
+                self.weight = torch.nn.Parameter(v, requires_grad=False)
+            elif k[prefix_len:] == "bias" and v is not None:
+                self.bias = torch.nn.Parameter(v, requires_grad=False)
             else:
                 missing_keys.append(k)
 
-    def _save_to_state_dict(self, destination, prefix, keep_vars):
+    def ggml_save_to_state_dict(self, destination, prefix, keep_vars):
         # This is a fake state dict for vram estimation
-        if self.weight is not None:
-            weight = torch.zeros_like(self.weight, device=torch.device("meta"))
-            destination[f"{prefix}weight"] = weight
+        weight = torch.zeros_like(self.weight, device=torch.device("meta"))
+        destination[prefix + "weight"] = weight
         if self.bias is not None:
             bias = torch.zeros_like(self.bias, device=torch.device("meta"))
-            destination[f"{prefix}bias"] = bias
+            destination[prefix + "bias"] = bias
         return
 
         # This would return the actual state dict
-        weight, bias = self.get_weights()
-        if weight is not None:
-            destination[f"{prefix}weight"] = weight
+        destination[prefix + "weight"] = self.get_weight(self.weight)
         if bias is not None:
-            destination[f"{prefix}bias"] = weight
-
-    def _apply(self, fn):
-        if self.weight is not None:
-            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
+            destination[prefix + "bias"] = self.get_weight(self.bias)
 
     def get_weight(self, tensor, dtype):
         if tensor is None:
@@ -128,9 +111,8 @@ class GGMLLayer(torch.nn.Module):
         # 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, key in getattr(tensor, "patches", []):
-            patch_list += move_patch_to_cuda(patches, device)
+            patch_list += move_patch_to_device(patches, device)
 
         # dequantize tensor while patches load
         weight = dequantize_tensor(tensor, dtype, self.dequant_dtype)
@@ -145,32 +127,75 @@ class GGMLLayer(torch.nn.Module):
                 weight = function(patch_list, weight, key, patch_dtype)
         return weight
 
-    def get_weights(self, dtype=torch.float16):
-        weight = self.get_weight(self.weight, dtype)
-        bias = self.get_weight(self.bias, dtype)
-        return (weight, bias)
+    def cast_bias_weight(s, input=None, dtype=None, device=None, bias_dtype=None):
+        if input is not None:
+            if dtype is None:
+                dtype = getattr(input, "dtype", torch.float32)
+            if bias_dtype is None:
+                bias_dtype = dtype
+            if device is None:
+                device = input.device
 
+        bias = None
+        non_blocking = comfy.model_management.device_supports_non_blocking(device)
+        if s.bias is not None:
+            bias = s.get_weight(s.bias.to(device), dtype)
+            bias = comfy.ops.cast_to(bias, bias_dtype, device, non_blocking=non_blocking, copy=False)
+
+        weight = s.get_weight(s.weight.to(device), dtype)
+        weight = comfy.ops.cast_to(weight, dtype, device, non_blocking=non_blocking, copy=False)
+        return weight, bias
+
+    def forward_comfy_cast_weights(self, input, *args, **kwargs):
+        if self.is_ggml_quantized():
+            return self.forward_ggml_cast_weights(input, *args, **kwargs)
+        return super().forward_comfy_cast_weights(input, *args, **kwargs)
+
+    def forward_ggml_cast_weights(self, input):
+        raise NotImplementedError
 
 class GGMLOps(comfy.ops.manual_cast):
     """
     Dequantize weights on the fly before doing the compute
     """
-    class Linear(GGMLLayer):
-        _forward_operation = staticmethod(torch.nn.functional.linear)
+    class Linear(GGMLLayer, comfy.ops.manual_cast.Linear):
+        def forward_ggml_cast_weights(self, input):
+            weight, bias = self.cast_bias_weight(input)
+            return torch.nn.functional.linear(input, weight, bias)
 
-    class Conv2d(GGMLLayer):
-        def __init__(self, *args, device=None, dtype=None, **kwargs):
-            super().__init__()
-            _ = kwargs.pop("kernel_size", None)
-            self._forward_operation = partial(staticmethod(torch.nn.functional.conv2d), **kwargs)
+    class Conv2d(GGMLLayer, comfy.ops.manual_cast.Conv2d):
+        def forward_ggml_cast_weights(self, input):
+            weight, bias = self.cast_bias_weight(input)
+            return self._conv_forward(input, weight, bias)
 
+    class Embedding(GGMLLayer, comfy.ops.manual_cast.Embedding):
+        def forward_ggml_cast_weights(self, input, out_dtype=None):
+            output_dtype = out_dtype
+            if self.weight.dtype == torch.float16 or self.weight.dtype == torch.bfloat16:
+                out_dtype = None
+            weight, _bias = self.cast_bias_weight(self, device=input.device, dtype=out_dtype)
+            return torch.nn.functional.embedding(
+                input, weight, self.padding_idx, self.max_norm, self.norm_type, self.scale_grad_by_freq, self.sparse
+            ).to(dtype=output_dtype)
 
-def move_patch_to_cuda(item, device):
+    class LayerNorm(GGMLLayer, comfy.ops.manual_cast.LayerNorm):
+        def forward_ggml_cast_weights(self, input):
+            if self.weight is None:
+                return super().forward_comfy_cast_weights(input)
+            weight, bias = self.cast_bias_weight(input)
+            return torch.nn.functional.layer_norm(input, self.normalized_shape, weight, bias, self.eps)
+
+    class GroupNorm(GGMLLayer, comfy.ops.manual_cast.GroupNorm):
+        def forward_ggml_cast_weights(self, input):
+            weight, bias = self.cast_bias_weight(input)
+            return torch.nn.functional.group_norm(input, self.num_groups, weight, bias, self.eps)
+
+def move_patch_to_device(item, device):
     if isinstance(item, torch.Tensor):
         return item.to(device, non_blocking=True)
     elif isinstance(item, tuple):
-        return tuple(move_patch_to_cuda(x, device) for x in item)
+        return tuple(move_patch_to_device(x, device) for x in item)
     elif isinstance(item, list):
-        return [move_patch_to_cuda(x, device) for x in item]
+        return [move_patch_to_device(x, device) for x in item]
     else:
         return item