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")