Refactor tensor type checking a bit more

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

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