Revert tensor construction change in dequantize_tensor

Eliminate GGMLOpsLayer, include functionality in GGMLLayer
This commit is contained in:
blepping
2024-08-25 10:35:39 -06:00
parent f3cf802dd6
commit 7413de86e6
2 changed files with 24 additions and 28 deletions
+2 -2
View File
@@ -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
+22 -26
View File
@@ -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)