Revert tensor construction change in dequantize_tensor
Eliminate GGMLOpsLayer, include functionality in GGMLLayer
This commit is contained in:
+2
-2
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user