From 7413de86e6ea4b63908625143525f10ca1a468dc Mon Sep 17 00:00:00 2001
From: blepping <blepping@tuta.io>
Date: Sun, 25 Aug 2024 10:35:39 -0600
Subject: [PATCH] Revert tensor construction change in dequantize_tensor

Eliminate GGMLOpsLayer, include functionality in GGMLLayer
---
 dequant.py |  4 ++--
 ops.py     | 48 ++++++++++++++++++++++--------------------------
 2 files changed, 24 insertions(+), 28 deletions(-)

diff --git a/dequant.py b/dequant.py
index cd9a8b7..1750cd0 100644
--- a/dequant.py
+++ b/dequant.py
@@ -4,11 +4,11 @@ import torch
 from tqdm import tqdm
 
 def dequantize_tensor(tensor, dtype=None, dequant_dtype=None):
-    data = tensor.data.clone().detach()
+    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]:
+    if qtype in (None, gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16):
         return data.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 1f9f666..2d67e23 100644
--- a/ops.py
+++ b/ops.py
@@ -56,6 +56,7 @@ class GGMLLayer(torch.nn.Module):
     """
     This (should) be responsible for de-quantizing on the fly
     """
+    comfy_cast_weights = True
     dequant_dtype = None
     patch_dtype = None
 
@@ -64,6 +65,24 @@ class GGMLLayer(torch.nn.Module):
         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)
+        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":
@@ -131,40 +150,17 @@ class GGMLLayer(torch.nn.Module):
         bias = self.get_weight(self.bias, dtype)
         return (weight, bias)
 
-class GGMLOpsLayer(GGMLLayer):
-    comfy_cast_weights = True
-
-    def __init__(self, *args, device=None, dtype=None, **kwargs):
-        super().__init__(device=device, dtype=dtype)
-
-    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)
-        del weight, bias
-
-        if device:
-            self.to(device)
-        return x
 
 class GGMLOps(comfy.ops.manual_cast):
     """
     Dequantize weights on the fly before doing the compute
     """
-    class Linear(GGMLOpsLayer):
+    class Linear(GGMLLayer):
         _forward_operation = staticmethod(torch.nn.functional.linear)
 
-    class Conv2d(GGMLOpsLayer):
+    class Conv2d(GGMLLayer):
         def __init__(self, *args, device=None, dtype=None, **kwargs):
-            super().__init__(device=device, dtype=dtype)
+            super().__init__()
             _ = kwargs.pop("kernel_size", None)
             self._forward_operation = partial(staticmethod(torch.nn.functional.conv2d), **kwargs)