From 6ae0c733b0edb283bc5a80ee7d679254eb230945 Mon Sep 17 00:00:00 2001
From: blepping <blepping@tuta.io>
Date: Wed, 28 Aug 2024 10:34:58 -0600
Subject: [PATCH 01/19] Support loading u-net from stable-diffusion.cpp GGUF
 files

Small cleanups
---
 nodes.py | 57 ++++++++++++++++++++++++++++++++++++++++----------------
 ops.py   | 36 ++++++++++++++++++++++++++++++-----
 2 files changed, 72 insertions(+), 21 deletions(-)

diff --git a/nodes.py b/nodes.py
index 0e61662..84c5def 100644
--- a/nodes.py
+++ b/nodes.py
@@ -32,27 +32,50 @@ 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.", require_prefix=False):
     """
     Read state dict as fake tensors
     """
     reader = gguf.GGUFReader(path)
     sd = {}
     dt = {}
+    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"}:
+            raise ValueError(f"Unexpected architecture type in GGUF file, expected one of flux, sd1, sdxl but got {arch_str!r}")
+    else:
+        arch_str = None
+    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)
+    else:
+        has_prefix = False
+    if require_prefix and not has_prefix:
+        raise RuntimeError(f"Loader requires prefix {handle_prefix!r} which does not exist in the model")
     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:]
+        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}:
+        if tensor.tensor_type in {gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16}:
             torch_tensor = torch_tensor.view(*shape)
-        sd[tensor_name] = GGMLTensor(
+        sd[sd_key] = GGMLTensor(
             torch_tensor,
             tensor_type = tensor.tensor_type,
-            tensor_shape = shape
+            tensor_shape = shape,
+            ggml_tensor_name = tensor_name,
         )
-        dt[str(tensor.tensor_type)] = dt.get(str(tensor.tensor_type), 0) + 1
+        dt[tensor_type_str] = dt.get(tensor_type_str, 0) + 1
 
     # sanity check debug print
     print("\nggml_sd_loader:")
@@ -172,14 +195,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
@@ -283,11 +306,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"), ),
             }
         }
 
@@ -296,18 +320,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,
             }
         }
 
@@ -317,7 +342,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)),)
 
diff --git a/ops.py b/ops.py
index 88ca894..21b4d36 100644
--- a/ops.py
+++ b/ops.py
@@ -9,13 +9,14 @@ class GGMLTensor(torch.Tensor):
     """
     Main tensor-like class for storing quantized weights
     """
-    def __init__(self, *args, tensor_type, tensor_shape, patches=[], **kwargs):
+    def __init__(self, *args, tensor_type, tensor_shape, patches=[], ggml_tensor_name=None, **kwargs):
         super().__init__()
         self.tensor_type = tensor_type
         self.tensor_shape = tensor_shape
         self.patches = patches
+        self.ggml_tensor_name = ggml_tensor_name
 
-    def __new__(cls, *args, tensor_type, tensor_shape, patches=[], **kwargs):
+    def __new__(cls, *args, tensor_type, tensor_shape, patches=[], ggml_tensor_name=None, **kwargs):
         return super().__new__(cls, *args, **kwargs)
 
     def to(self, *args, **kwargs):
@@ -23,6 +24,7 @@ class GGMLTensor(torch.Tensor):
         new.tensor_type = getattr(self, "tensor_type", None)
         new.tensor_shape = getattr(self, "tensor_shape", new.data.shape)
         new.patches = getattr(self, "patches", []).copy()
+        new.ggml_tensor_name = getattr(self, "ggml_tensor_name", None)
         return new
 
     def clone(self, *args, **kwargs):
@@ -44,6 +46,7 @@ class GGMLTensor(torch.Tensor):
         new.tensor_type = getattr(self, "tensor_type", None)
         new.tensor_shape = getattr(self, "tensor_shape", new.data.shape)
         new.patches = getattr(self, "patches", []).copy()
+        new.ggml_tensor_name = getattr(self, "ggml_tensor_name", None)
         return new
 
     @property
@@ -76,7 +79,7 @@ class GGMLLayer(torch.nn.Module):
             self.to(x.device)
 
         weight, bias = self.get_weights(x.dtype)
-        x = self._forward_operation(x, weight, bias=bias)
+        x = self._forward_operation(input=x, weight=weight, bias=bias)
         del weight, bias
 
         if device:
@@ -84,10 +87,11 @@ class GGMLLayer(torch.nn.Module):
         return x
 
     def _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":
+            if k[prefix_len:] == "weight":
                 self.weight = v
-            elif k[len(prefix):] == "bias":
+            elif k[prefix_len:] == "bias":
                 self.bias = v
             else:
                 missing_keys.append(k)
@@ -164,6 +168,28 @@ class GGMLOps(comfy.ops.manual_cast):
             _ = kwargs.pop("kernel_size", None)
             self._forward_operation = partial(staticmethod(torch.nn.functional.conv2d), **kwargs)
 
+    class GroupNorm(GGMLLayer):
+        def __init__(self, *args, device=None, dtype=None, **kwargs):
+            super().__init__()
+            num_groups = kwargs.get("num_groups")
+            if num_groups is None:
+                if len(args) == 0 or not isinstance(args[0], int):
+                    raise ValueError("Couldn't get num_groups in GroupNorm: missing or bad type")
+                num_groups = args[0]
+            eps = kwargs.get("eps", 1e-05)
+            self._forward_operation = partial(staticmethod(torch.nn.functional.group_norm), num_groups=num_groups, eps=eps)
+
+    class LayerNorm(GGMLLayer):
+        def __init__(self, *args, device=None, dtype=None, **kwargs):
+            super().__init__()
+            normalized_shape = kwargs.get("normalized_shape")
+            if normalized_shape is None:
+                if len(args) == 0 or not all(isinstance(v, int) for v in args):
+                    raise ValueError("Couldn't get normalized_shape in LayerNorm: missing or bad type")
+                normalized_shape = tuple(args)
+            eps = kwargs.get("eps", 1e-05)
+            self._forward_operation = partial(staticmethod(torch.nn.functional.layer_norm), normalized_shape=normalized_shape, eps=eps)
+
 
 def move_patch_to_cuda(item, device):
     if isinstance(item, torch.Tensor):

From f7d8fbe74052e139a3b20b4a4f702bda82fdc23e Mon Sep 17 00:00:00 2001
From: blepping <blepping@tuta.io>
Date: Wed, 28 Aug 2024 11:56:51 -0600
Subject: [PATCH 02/19] Workaround for stable-diffusion.cpp SDXL GGUF files

---
 nodes.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/nodes.py b/nodes.py
index 84c5def..9df2523 100644
--- a/nodes.py
+++ b/nodes.py
@@ -67,6 +67,11 @@ def gguf_sd_loader(path, handle_prefix="model.diffusion_model.", require_prefix=
         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)))
+            if arch_str is None and (tensor_name.endswith(".proj_in.weight") or tensor_name.endswith(".proj_out.weight")):
+                # Workaround for stable-diffusion.cpp SDXL detection.
+                # Impossible to land here for our own GGUF files as we set architecture.
+                while len(shape) > 2 and shape[-1] == 1:
+                    shape = shape[:-1]
         if tensor.tensor_type in {gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16}:
             torch_tensor = torch_tensor.view(*shape)
         sd[sd_key] = GGMLTensor(

From e814fc12642aa4d6ad3ae5323614c028ad1c2be1 Mon Sep 17 00:00:00 2001
From: blepping <blepping@tuta.io>
Date: Wed, 28 Aug 2024 12:34:08 -0600
Subject: [PATCH 03/19] Revert ggml_tensor_name addition

---
 nodes.py | 3 +--
 ops.py   | 7 ++-----
 2 files changed, 3 insertions(+), 7 deletions(-)

diff --git a/nodes.py b/nodes.py
index 9df2523..d702e0c 100644
--- a/nodes.py
+++ b/nodes.py
@@ -77,8 +77,7 @@ def gguf_sd_loader(path, handle_prefix="model.diffusion_model.", require_prefix=
         sd[sd_key] = GGMLTensor(
             torch_tensor,
             tensor_type = tensor.tensor_type,
-            tensor_shape = shape,
-            ggml_tensor_name = tensor_name,
+            tensor_shape = shape
         )
         dt[tensor_type_str] = dt.get(tensor_type_str, 0) + 1
 
diff --git a/ops.py b/ops.py
index 21b4d36..8c12de7 100644
--- a/ops.py
+++ b/ops.py
@@ -9,14 +9,13 @@ class GGMLTensor(torch.Tensor):
     """
     Main tensor-like class for storing quantized weights
     """
-    def __init__(self, *args, tensor_type, tensor_shape, patches=[], ggml_tensor_name=None, **kwargs):
+    def __init__(self, *args, tensor_type, tensor_shape, patches=[], **kwargs):
         super().__init__()
         self.tensor_type = tensor_type
         self.tensor_shape = tensor_shape
         self.patches = patches
-        self.ggml_tensor_name = ggml_tensor_name
 
-    def __new__(cls, *args, tensor_type, tensor_shape, patches=[], ggml_tensor_name=None, **kwargs):
+    def __new__(cls, *args, tensor_type, tensor_shape, patches=[], **kwargs):
         return super().__new__(cls, *args, **kwargs)
 
     def to(self, *args, **kwargs):
@@ -24,7 +23,6 @@ class GGMLTensor(torch.Tensor):
         new.tensor_type = getattr(self, "tensor_type", None)
         new.tensor_shape = getattr(self, "tensor_shape", new.data.shape)
         new.patches = getattr(self, "patches", []).copy()
-        new.ggml_tensor_name = getattr(self, "ggml_tensor_name", None)
         return new
 
     def clone(self, *args, **kwargs):
@@ -46,7 +44,6 @@ class GGMLTensor(torch.Tensor):
         new.tensor_type = getattr(self, "tensor_type", None)
         new.tensor_shape = getattr(self, "tensor_shape", new.data.shape)
         new.patches = getattr(self, "patches", []).copy()
-        new.ggml_tensor_name = getattr(self, "ggml_tensor_name", None)
         return new
 
     @property

From b88cb9a344d5f8601dbb9e9af13c9d9c4b466112 Mon Sep 17 00:00:00 2001
From: blepping <blepping@tuta.io>
Date: Wed, 28 Aug 2024 12:50:39 -0600
Subject: [PATCH 04/19] Remove require_prefix in gguf_sd_loader

---
 nodes.py | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/nodes.py b/nodes.py
index d702e0c..609716e 100644
--- a/nodes.py
+++ b/nodes.py
@@ -32,7 +32,7 @@ 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, handle_prefix="model.diffusion_model.", require_prefix=False):
+def gguf_sd_loader(path, handle_prefix="model.diffusion_model."):
     """
     Read state dict as fake tensors
     """
@@ -54,8 +54,6 @@ def gguf_sd_loader(path, handle_prefix="model.diffusion_model.", require_prefix=
         has_prefix = any(s.startswith(handle_prefix) for s in tensor_names)
     else:
         has_prefix = False
-    if require_prefix and not has_prefix:
-        raise RuntimeError(f"Loader requires prefix {handle_prefix!r} which does not exist in the model")
     for tensor in reader.tensors:
         sd_key = tensor_name = tensor.name
         if has_prefix:

From 826bd1804cc67e59bb23a9cbcfdf2ec65304b984 Mon Sep 17 00:00:00 2001
From: City <125218114+city96@users.noreply.github.com>
Date: Fri, 30 Aug 2024 03:25:57 +0200
Subject: [PATCH 05/19] Make ops less jank

Also possibly faster. Ignore the massively increased system memory usage.
---
 ops.py | 81 +++++++++++++++++++++-------------------------------------
 1 file changed, 29 insertions(+), 52 deletions(-)

diff --git a/ops.py b/ops.py
index 88ca894..d78f982 100644
--- a/ops.py
+++ b/ops.py
@@ -1,6 +1,6 @@
 # (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
@@ -60,35 +60,12 @@ class GGMLLayer(torch.nn.Module):
     dequant_dtype = None
     patch_dtype = None
 
-    def __init__(self, *args, **kwargs):
-        super().__init__()
-        self.weight = GGMLTensor(1, tensor_type=None, tensor_shape=None)
-        self.bias = None
-
-    def _forward_operation(self, x, weight, bias):
-        raise NotImplementedError
-
-    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 = 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):
         for k,v in state_dict.items():
             if k[len(prefix):] == "weight":
-                self.weight = v
-            elif k[len(prefix):] == "bias":
-                self.bias = v
+                self.weight = torch.nn.Parameter(v, requires_grad=False)
+            elif k[len(prefix):] == "bias" and v is not None:
+                self.bias = torch.nn.Parameter(v, requires_grad=False)
             else:
                 missing_keys.append(k)
 
@@ -109,18 +86,6 @@ class GGMLLayer(torch.nn.Module):
         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
-
     def get_weight(self, tensor, dtype):
         if tensor is None:
             return
@@ -128,7 +93,6 @@ 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)
 
@@ -145,25 +109,38 @@ 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 = input.dtype
+            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) if hasattr(s, "get_weight") else s.bias
+            bias = comfy.ops.cast_to(bias, bias_dtype, device, non_blocking=non_blocking, copy=False)
+
+        weight = s.get_weight(s.weight.to(device), dtype) if hasattr(s, "get_weight") else s.weight
+        weight = comfy.ops.cast_to(weight, dtype, device, non_blocking=non_blocking, copy=False)
+        return weight, bias
 
 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 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 Linear(GGMLLayer, comfy.ops.manual_cast.Linear):
+        def forward_comfy_cast_weights(self, input):
+            weight, bias = self.cast_bias_weight(input)
+            return torch.nn.functional.linear(input, weight, bias)
 
+    class Conv2d(GGMLLayer, comfy.ops.manual_cast.Conv2d):
+        def forward_comfy_cast_weights(self, input):
+            weight, bias = self.cast_bias_weight(input)
+            return self._conv_forward(input, weight, bias)
 
 def move_patch_to_cuda(item, device):
     if isinstance(item, torch.Tensor):

From f40b149bc5117582828c399a382d60c9781896e6 Mon Sep 17 00:00:00 2001
From: City <125218114+city96@users.noreply.github.com>
Date: Fri, 30 Aug 2024 03:39:57 +0200
Subject: [PATCH 06/19] Add BF16 support

---
 dequant.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/dequant.py b/dequant.py
index 1750cd0..454bd8c 100644
--- a/dequant.py
+++ b/dequant.py
@@ -45,6 +45,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 +226,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,

From ec24e75d89402a2f4d113d01f0e255a006632b31 Mon Sep 17 00:00:00 2001
From: City <125218114+city96@users.noreply.github.com>
Date: Fri, 30 Aug 2024 03:40:15 +0200
Subject: [PATCH 07/19] Remove unused check

---
 ops.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/ops.py b/ops.py
index d78f982..dbdaed9 100644
--- a/ops.py
+++ b/ops.py
@@ -121,10 +121,10 @@ class GGMLLayer(torch.nn.Module):
         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) if hasattr(s, "get_weight") else s.bias
+            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) if hasattr(s, "get_weight") else s.weight
+        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
 

From c231239fdfb6e3c484f091c6ea004f767ee2ea1f Mon Sep 17 00:00:00 2001
From: City <125218114+city96@users.noreply.github.com>
Date: Fri, 30 Aug 2024 03:47:52 +0200
Subject: [PATCH 08/19] cleanup

Perf difference from `torch.tensor(tensor.data)` doesn't seem realistically measurable and the warning is annoying.
---
 dequant.py | 11 +++++------
 ops.py     | 13 +++++--------
 2 files changed, 10 insertions(+), 14 deletions(-)

diff --git a/dequant.py b/dequant.py
index 454bd8c..3688d7f 100644
--- a/dequant.py
+++ b/dequant.py
@@ -4,20 +4,19 @@ 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)
+    oshape = getattr(tensor, "tensor_shape", tensor.shape)
 
     if qtype in (None, gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16):
-        return data.to(dtype)
+        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, 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):
     """
diff --git a/ops.py b/ops.py
index dbdaed9..43f6299 100644
--- a/ops.py
+++ b/ops.py
@@ -71,20 +71,17 @@ class GGMLLayer(torch.nn.Module):
 
     def _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
+            destination[prefix + "bias"] = self.get_weight(self.bias)
 
     def get_weight(self, tensor, dtype):
         if tensor is None:

From 6618801f1c1ee6c4203e4a76b87d17f3a2a3d3eb Mon Sep 17 00:00:00 2001
From: City <125218114+city96@users.noreply.github.com>
Date: Fri, 30 Aug 2024 06:25:19 +0200
Subject: [PATCH 09/19] Don't dequantize token embed on load

---
 nodes.py |  7 -------
 ops.py   | 12 +++++++++++-
 2 files changed, 11 insertions(+), 8 deletions(-)

diff --git a/nodes.py b/nodes.py
index 70823f4..9b7bede 100644
--- a/nodes.py
+++ b/nodes.py
@@ -78,10 +78,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 +86,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
 
diff --git a/ops.py b/ops.py
index 43f6299..bc17c3a 100644
--- a/ops.py
+++ b/ops.py
@@ -109,7 +109,7 @@ class GGMLLayer(torch.nn.Module):
     def cast_bias_weight(s, input=None, dtype=None, device=None, bias_dtype=None):
         if input is not None:
             if dtype is None:
-                dtype = input.dtype
+                dtype = getattr(input, "dtype", torch.float32)
             if bias_dtype is None:
                 bias_dtype = dtype
             if device is None:
@@ -139,6 +139,16 @@ class GGMLOps(comfy.ops.manual_cast):
             weight, bias = self.cast_bias_weight(input)
             return self._conv_forward(input, weight, bias)
 
+    class Embedding(GGMLLayer, comfy.ops.manual_cast.Embedding):
+        def forward_comfy_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):
     if isinstance(item, torch.Tensor):
         return item.to(device, non_blocking=True)

From a0cb97ecebc0b9a1ebeaac77cd201a30948da402 Mon Sep 17 00:00:00 2001
From: blepping <blepping@tuta.io>
Date: Sun, 1 Sep 2024 05:25:32 -0600
Subject: [PATCH 10/19] Add LayerNorm and GroupNorm ops

Delegate to superclass methods whenever possible
---
 dequant.py |  2 +-
 ops.py     | 57 ++++++++++++++++++++++++++++++++++++++++++++++--------
 2 files changed, 50 insertions(+), 9 deletions(-)

diff --git a/dequant.py b/dequant.py
index 3688d7f..5bc775b 100644
--- a/dequant.py
+++ b/dequant.py
@@ -11,7 +11,7 @@ def dequantize_tensor(tensor, dtype=None, dequant_dtype=None):
         return tensor.to(dtype)
     elif qtype in dequantize_functions:
         dequant_dtype = dtype if dequant_dtype == "target" else dequant_dtype
-        return dequantize(tensor, 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}")
diff --git a/ops.py b/ops.py
index a6f47c1..70a0111 100644
--- a/ops.py
+++ b/ops.py
@@ -59,18 +59,39 @@ class GGMLLayer(torch.nn.Module):
     comfy_cast_weights = True
     dequant_dtype = None
     patch_dtype = None
+    ggml_delegate_types = {None, gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16}
 
-    def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs):
+    def ggml_check_delegate(self, *, weight=None, bias=None):
+        if weight is None:
+            weight = self.weight
+        if bias is None:
+            bias = self.bias
+        weight_type = getattr(weight, "tensor_type", None)
+        bias_type = None if bias is None else getattr(bias, "tensor_type", None)
+        return weight_type in self.ggml_delegate_types and bias_type in self.ggml_delegate_types
+
+    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.ggml_check_delegate(weight=weight, bias=bias):
+            return super()._load_from_state_dict(state_dict, prefix, *args, **kwargs)
+        return self.ggml_load_from_state_dict(state_dict, prefix, *args, **kwargs)
+
+    def _save_to_state_dict(self, *args, **kwargs):
+        if self.ggml_check_delegate():
+            return super()._save_to_state_dict(*args, **kwargs)
+        return self.ggml_save_to_state_dict(*args, **kwargs)
+
+    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[prefix_len] == "weight":
+            if k[prefix_len:] == "weight":
                 self.weight = torch.nn.Parameter(v, requires_grad=False)
-            elif k[prefix_len] == "bias" and v is not None:
+            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
         weight = torch.zeros_like(self.weight, device=torch.device("meta"))
         destination[prefix + "weight"] = weight
@@ -126,30 +147,50 @@ class GGMLLayer(torch.nn.Module):
         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.ggml_check_delegate():
+            return super().forward_comfy_cast_weights(input, *args, **kwargs)
+        return self.ggml_forward(input, *args, **kwargs)
+
+    def ggml_forward(self, input):
+        raise NotImplementedError
+
 class GGMLOps(comfy.ops.manual_cast):
     """
     Dequantize weights on the fly before doing the compute
     """
     class Linear(GGMLLayer, comfy.ops.manual_cast.Linear):
-        def forward_comfy_cast_weights(self, input):
+        def ggml_forward(self, input):
             weight, bias = self.cast_bias_weight(input)
             return torch.nn.functional.linear(input, weight, bias)
 
     class Conv2d(GGMLLayer, comfy.ops.manual_cast.Conv2d):
-        def forward_comfy_cast_weights(self, input):
+        def ggml_forward(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_comfy_cast_weights(self, input, out_dtype=None):
+        def ggml_forward(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)
+            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)
 
+    class LayerNorm(GGMLLayer, comfy.ops.manual_cast.LayerNorm):
+        def ggml_forward(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 ggml_forward(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_cuda(item, device):
     if isinstance(item, torch.Tensor):
         return item.to(device, non_blocking=True)

From 3bad1b4e04e839fd6ec57f8ea0b2a657e4936034 Mon Sep 17 00:00:00 2001
From: blepping <blepping@tuta.io>
Date: Sun, 1 Sep 2024 06:46:50 -0600
Subject: [PATCH 11/19] Fix loading GGUF t5xxl

Can't use squeeze proj workaround for stable-diffusion.cpp SD15
---
 nodes.py | 14 ++++++++++----
 1 file changed, 10 insertions(+), 4 deletions(-)

diff --git a/nodes.py b/nodes.py
index d095de2..df5678c 100644
--- a/nodes.py
+++ b/nodes.py
@@ -11,7 +11,7 @@ import comfy.model_patcher
 import folder_paths
 
 from .ops import GGMLTensor, GGMLOps, move_patch_to_cuda
-from .dequant import dequantize_tensor
+from .tools.convert import detect_arch
 
 # Add a custom keys for files ending in .gguf
 if "unet_gguf" not in folder_paths.folder_names_and_paths:
@@ -44,8 +44,8 @@ def gguf_sd_loader(path, handle_prefix="model.diffusion_model."):
         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"}:
-            raise ValueError(f"Unexpected architecture type in GGUF file, expected one of flux, sd1, sdxl but got {arch_str!r}")
+        if arch_str not in {"flux", "sd1", "sdxl", "t5encoder"}:
+            raise ValueError(f"Unexpected architecture type in GGUF file, expected one of flux, sd1, sdxl, t5encoder but got {arch_str!r}")
     else:
         arch_str = None
     if handle_prefix is not None:
@@ -54,18 +54,24 @@ def gguf_sd_loader(path, handle_prefix="model.diffusion_model."):
         has_prefix = any(s.startswith(handle_prefix) for s in tensor_names)
     else:
         has_prefix = False
+    tensors = []
     for tensor in reader.tensors:
         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))
+    detected_arch = detect_arch(set(val[0] for val in tensors)) if arch_str is None else None
+    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)))
-            if arch_str is None and (tensor_name.endswith(".proj_in.weight") or tensor_name.endswith(".proj_out.weight")):
+            if arch_str is None and detected_arch == "sdxl" and \
+                (tensor_name.endswith(".proj_in.weight") or tensor_name.endswith(".proj_out.weight")):
                 # Workaround for stable-diffusion.cpp SDXL detection.
                 # Impossible to land here for our own GGUF files as we set architecture.
                 while len(shape) > 2 and shape[-1] == 1:

From 43427eff52db19265368d0c17d7cffb0ac11c0fc Mon Sep 17 00:00:00 2001
From: City <125218114+city96@users.noreply.github.com>
Date: Mon, 2 Sep 2024 01:37:16 +0200
Subject: [PATCH 12/19] Better patch clear logic

Reduces RAM overhead
---
 nodes.py | 28 +++++++++++++++++++---------
 1 file changed, 19 insertions(+), 9 deletions(-)

diff --git a/nodes.py b/nodes.py
index 9b7bede..58934b0 100644
--- a/nodes.py
+++ b/nodes.py
@@ -98,9 +98,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
@@ -110,16 +107,17 @@ class GGUFModelPatcher(comfy.model_patcher.ModelPatcher):
         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
-
+            out_weight = weight.to(device_to)
             if self.patch_on_device:
                 patches = move_patch_to_cuda(patches, self.load_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:
@@ -133,6 +131,18 @@ 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():
+                qtype = getattr(p, "tensor_type", None) 
+                patches = getattr(p, "patches", [])
+                if qtype in (None, gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16):
+                    continue
+                if len(patches) > 0:
+                    p.patches = []
+
+        return super().unpatch_model(device_to=device_to, 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)

From c7b902bf9a7f8a020a370249d56f9b339c447c65 Mon Sep 17 00:00:00 2001
From: City <125218114+city96@users.noreply.github.com>
Date: Mon, 2 Sep 2024 02:00:10 +0200
Subject: [PATCH 13/19] Patch load/unload fix

May break other things. For some reason T5 decided to run on the CPU with this change, but only like once and I can't for the life of me reproduce it.
---
 nodes.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/nodes.py b/nodes.py
index 58934b0..062aafe 100644
--- a/nodes.py
+++ b/nodes.py
@@ -140,8 +140,7 @@ class GGUFModelPatcher(comfy.model_patcher.ModelPatcher):
                     continue
                 if len(patches) > 0:
                     p.patches = []
-
-        return super().unpatch_model(device_to=device_to, unpatch_weights=unpatch_weights)
+        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

From e19fe8e898231b5ec35d385e2ea3a9373f85bc6e Mon Sep 17 00:00:00 2001
From: blepping <blepping@tuta.io>
Date: Sun, 1 Sep 2024 18:04:31 -0600
Subject: [PATCH 14/19] Rename ggml_check_delegate to is_ggml_quantized,
 refactor accordingly

---
 ops.py | 24 ++++++++++++------------
 1 file changed, 12 insertions(+), 12 deletions(-)

diff --git a/ops.py b/ops.py
index 70a0111..930b4da 100644
--- a/ops.py
+++ b/ops.py
@@ -59,27 +59,27 @@ class GGMLLayer(torch.nn.Module):
     comfy_cast_weights = True
     dequant_dtype = None
     patch_dtype = None
-    ggml_delegate_types = {None, gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16}
+    torch_compatible_tensor_types = {None, gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16}
 
-    def ggml_check_delegate(self, *, weight=None, 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
         weight_type = getattr(weight, "tensor_type", None)
         bias_type = None if bias is None else getattr(bias, "tensor_type", None)
-        return weight_type in self.ggml_delegate_types and bias_type in self.ggml_delegate_types
+        return any(t not in self.torch_compatible_tensor_types for t in (weight_type, bias_type))
 
     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.ggml_check_delegate(weight=weight, bias=bias):
-            return super()._load_from_state_dict(state_dict, prefix, *args, **kwargs)
-        return self.ggml_load_from_state_dict(state_dict, prefix, *args, **kwargs)
+        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 _save_to_state_dict(self, *args, **kwargs):
-        if self.ggml_check_delegate():
-            return super()._save_to_state_dict(*args, **kwargs)
-        return self.ggml_save_to_state_dict(*args, **kwargs)
+        if self.is_ggml_quantized():
+            return self.ggml_save_to_state_dict(*args, **kwargs)
+        return super()._save_to_state_dict(*args, **kwargs)
 
     def ggml_load_from_state_dict(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs):
         prefix_len = len(prefix)
@@ -148,9 +148,9 @@ class GGMLLayer(torch.nn.Module):
         return weight, bias
 
     def forward_comfy_cast_weights(self, input, *args, **kwargs):
-        if self.ggml_check_delegate():
-            return super().forward_comfy_cast_weights(input, *args, **kwargs)
-        return self.ggml_forward(input, *args, **kwargs)
+        if self.is_ggml_quantized():
+            return self.ggml_forward(input, *args, **kwargs)
+        return super().forward_comfy_cast_weights(input, *args, **kwargs)
 
     def ggml_forward(self, input):
         raise NotImplementedError

From ba8ce44b4c3e0ce8c5a07529983e6fb58b9d1f1e Mon Sep 17 00:00:00 2001
From: City <125218114+city96@users.noreply.github.com>
Date: Mon, 2 Sep 2024 02:22:32 +0200
Subject: [PATCH 15/19] Rename move_patch function

---
 nodes.py | 5 ++---
 ops.py   | 8 ++++----
 2 files changed, 6 insertions(+), 7 deletions(-)

diff --git a/nodes.py b/nodes.py
index 062aafe..094d8e9 100644
--- a/nodes.py
+++ b/nodes.py
@@ -10,7 +10,7 @@ import comfy.model_management
 import comfy.model_patcher
 import folder_paths
 
-from .ops import GGMLTensor, GGMLOps, move_patch_to_cuda
+from .ops import GGMLTensor, GGMLOps, move_patch_to_device
 from .dequant import dequantize_tensor
 
 # Add a custom keys for files ending in .gguf
@@ -108,8 +108,7 @@ class GGUFModelPatcher(comfy.model_patcher.ModelPatcher):
         qtype = getattr(weight, "tensor_type", None)
         if qtype not in (None, gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16):
             out_weight = weight.to(device_to)
-            if self.patch_on_device:
-                patches = move_patch_to_cuda(patches, self.load_device)
+            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
diff --git a/ops.py b/ops.py
index bc17c3a..82a3e8e 100644
--- a/ops.py
+++ b/ops.py
@@ -91,7 +91,7 @@ class GGMLLayer(torch.nn.Module):
         patch_list = []
         device = tensor.device
         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)
@@ -149,12 +149,12 @@ class GGMLOps(comfy.ops.manual_cast):
                 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):
+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

From 379175e7bf8b65019cdd11108bb882120a6f17df Mon Sep 17 00:00:00 2001
From: blepping <blepping@tuta.io>
Date: Sun, 1 Sep 2024 18:30:47 -0600
Subject: [PATCH 16/19] Refactor tensor type checking a bit more

---
 dequant.py | 11 ++++++++++-
 ops.py     |  6 ++----
 2 files changed, 12 insertions(+), 5 deletions(-)

diff --git a/dequant.py b/dequant.py
index 5bc775b..8232500 100644
--- a/dequant.py
+++ b/dequant.py
@@ -3,11 +3,20 @@ import gguf
 import torch
 from tqdm import tqdm
 
+
+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 (None, gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16):
+    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
diff --git a/ops.py b/ops.py
index 930b4da..1789b7f 100644
--- a/ops.py
+++ b/ops.py
@@ -3,7 +3,7 @@ import gguf
 import torch
 
 import comfy.ops
-from .dequant import dequantize_tensor
+from .dequant import dequantize_tensor, is_quantized
 
 class GGMLTensor(torch.Tensor):
     """
@@ -66,9 +66,7 @@ class GGMLLayer(torch.nn.Module):
             weight = self.weight
         if bias is None:
             bias = self.bias
-        weight_type = getattr(weight, "tensor_type", None)
-        bias_type = None if bias is None else getattr(bias, "tensor_type", None)
-        return any(t not in self.torch_compatible_tensor_types for t in (weight_type, bias_type))
+        return is_quantized(weight) or is_quantized(bias)
 
     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")

From df3378da6bebcde6181a5011e94391e9ceda7542 Mon Sep 17 00:00:00 2001
From: City <125218114+city96@users.noreply.github.com>
Date: Mon, 2 Sep 2024 03:17:16 +0200
Subject: [PATCH 17/19] Detect quant

---
 nodes.py | 9 ++++-----
 1 file changed, 4 insertions(+), 5 deletions(-)

diff --git a/nodes.py b/nodes.py
index ed4c242..e4ad05e 100644
--- a/nodes.py
+++ b/nodes.py
@@ -12,6 +12,7 @@ import folder_paths
 
 from .ops import GGMLTensor, GGMLOps, move_patch_to_device
 from .tools.convert import detect_arch
+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:
@@ -136,8 +137,7 @@ 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 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))
@@ -164,10 +164,9 @@ class GGUFModelPatcher(comfy.model_patcher.ModelPatcher):
     def unpatch_model(self, device_to=None, unpatch_weights=True):
         if unpatch_weights:
             for p in self.model.parameters():
-                qtype = getattr(p, "tensor_type", None)
-                patches = getattr(p, "patches", [])
-                if qtype in (None, gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16):
+                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)

From 35007f2659478359e4386befe31519e0c067e1a1 Mon Sep 17 00:00:00 2001
From: City <125218114+city96@users.noreply.github.com>
Date: Mon, 2 Sep 2024 03:44:07 +0200
Subject: [PATCH 18/19] Refactor model loading

Should make it more readable. Moved import under if to avoid changes to stuff in `tools` breaking normal models.
---
 nodes.py | 67 +++++++++++++++++++++++++++++++-------------------------
 1 file changed, 37 insertions(+), 30 deletions(-)

diff --git a/nodes.py b/nodes.py
index e4ad05e..e129988 100644
--- a/nodes.py
+++ b/nodes.py
@@ -11,7 +11,6 @@ import comfy.model_patcher
 import folder_paths
 
 from .ops import GGMLTensor, GGMLOps, move_patch_to_device
-from .tools.convert import detect_arch
 from .dequant import is_quantized, is_torch_compatible
 
 # Add a custom keys for files ending in .gguf
@@ -38,23 +37,14 @@ def gguf_sd_loader(path, handle_prefix="model.diffusion_model."):
     Read state dict as fake tensors
     """
     reader = gguf.GGUFReader(path)
-    sd = {}
-    dt = {}
-    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", "t5encoder"}:
-            raise ValueError(f"Unexpected architecture type in GGUF file, expected one of flux, sd1, sdxl, t5encoder but got {arch_str!r}")
-    else:
-        arch_str = None
+
+    # 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)
-    else:
-        has_prefix = False
+
     tensors = []
     for tensor in reader.tensors:
         sd_key = tensor_name = tensor.name
@@ -63,35 +53,52 @@ def gguf_sd_loader(path, handle_prefix="model.diffusion_model."):
                 continue
             sd_key = tensor_name[prefix_len:]
         tensors.append((sd_key, tensor))
-    detected_arch = detect_arch(set(val[0] for val in tensors)) if arch_str is None else None
+
+    # 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)))
-            if arch_str is None and detected_arch == "sdxl" and \
-                (tensor_name.endswith(".proj_in.weight") or tensor_name.endswith(".proj_out.weight")):
-                # Workaround for stable-diffusion.cpp SDXL detection.
-                # Impossible to land here for our own GGUF files as we set architecture.
-                while len(shape) > 2 and shape[-1] == 1:
-                    shape = shape[:-1]
+            # 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[sd_key] = GGMLTensor(
-            torch_tensor,
-            tensor_type = tensor.tensor_type,
-            tensor_shape = shape
-        )
-        dt[tensor_type_str] = dt.get(tensor_type_str, 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 = {

From 929d15454dd35d3e15eabdcc25d91dbb41f17f25 Mon Sep 17 00:00:00 2001
From: City <125218114+city96@users.noreply.github.com>
Date: Mon, 2 Sep 2024 03:58:26 +0200
Subject: [PATCH 19/19] Rename ggml_forward

Makes it clear that this is the same thing as comfy_cast_weights
---
 ops.py | 15 ++++++++-------
 1 file changed, 8 insertions(+), 7 deletions(-)

diff --git a/ops.py b/ops.py
index 6e106fe..62e797c 100644
--- a/ops.py
+++ b/ops.py
@@ -3,6 +3,7 @@ import gguf
 import torch
 
 import comfy.ops
+import comfy.model_management
 from .dequant import dequantize_tensor, is_quantized
 
 class GGMLTensor(torch.Tensor):
@@ -147,10 +148,10 @@ class GGMLLayer(torch.nn.Module):
 
     def forward_comfy_cast_weights(self, input, *args, **kwargs):
         if self.is_ggml_quantized():
-            return self.ggml_forward(input, *args, **kwargs)
+            return self.forward_ggml_cast_weights(input, *args, **kwargs)
         return super().forward_comfy_cast_weights(input, *args, **kwargs)
 
-    def ggml_forward(self, input):
+    def forward_ggml_cast_weights(self, input):
         raise NotImplementedError
 
 class GGMLOps(comfy.ops.manual_cast):
@@ -158,17 +159,17 @@ class GGMLOps(comfy.ops.manual_cast):
     Dequantize weights on the fly before doing the compute
     """
     class Linear(GGMLLayer, comfy.ops.manual_cast.Linear):
-        def ggml_forward(self, input):
+        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, comfy.ops.manual_cast.Conv2d):
-        def ggml_forward(self, input):
+        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 ggml_forward(self, input, out_dtype=None):
+        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
@@ -178,14 +179,14 @@ class GGMLOps(comfy.ops.manual_cast):
             ).to(dtype=output_dtype)
 
     class LayerNorm(GGMLLayer, comfy.ops.manual_cast.LayerNorm):
-        def ggml_forward(self, input):
+        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 ggml_forward(self, input):
+        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)