From 91bf0774dcea9c2f913defd01b2d35339ede4204 Mon Sep 17 00:00:00 2001 From: lspindler Date: Mon, 8 Sep 2025 13:01:36 +0200 Subject: [PATCH] Ensure backward compatabilty with 3rd party tools like GGUF --- comfy/ops.py | 40 ++++++++++++++++++++++++++++++---------- 1 file changed, 30 insertions(+), 10 deletions(-) diff --git a/comfy/ops.py b/comfy/ops.py index c2cf1d386..3ec2207eb 100644 --- a/comfy/ops.py +++ b/comfy/ops.py @@ -22,7 +22,8 @@ import comfy.model_management import comfy.float import comfy.rmsnorm import contextlib -from comfy.quant_tensor import Q_TYPES, tensor_quantizer, tensor_dequantizer, dynamic_tensor_quantizer, woq_fwd, quantized_fwd, get_quantizer_with_constraints +from comfy.quant_tensor import (Q_TYPES, tensor_quantizer, tensor_dequantizer, dynamic_tensor_quantizer, + woq_fwd, quantized_fwd, get_quantizer_with_constraints) import types import inspect @@ -99,7 +100,7 @@ class CastWeightBiasOp: use_dynamic_quantizer: bool = False class disable_weight_init: - class Linear(torch.nn.Module, CastWeightBiasOp): + class Linear(torch.nn.Linear, CastWeightBiasOp): def __init__( self, in_features: int, @@ -109,7 +110,7 @@ class disable_weight_init: dtype=None, ) -> None: factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() + torch.nn.Module.__init__(self) self.in_features = in_features self.out_features = out_features @@ -162,6 +163,9 @@ class disable_weight_init: state_dict[f"{prefix}scale_weight"] = scale_weight weight_dtype = _w.dtype + if weight_dtype not in Q_TYPES: + weight_dtype = self.compute_dtype + weight = torch.nn.Parameter( torch.empty((self.out_features, self.in_features), device=self.device, dtype=weight_dtype)) @@ -171,7 +175,7 @@ class disable_weight_init: scale_weight = state_dict.get(f"{prefix}scale_weight", None) if scale_weight is None: - logging.warning("Using quantized weights without a scale can result in low accuracy.") + logging.debug("Using quantized weights without a scale can result in low accuracy.") scale_weight = torch.ones(1) state_dict[f"{prefix}scale_weight"] = scale_weight self.register_buffer('scale_weight', scale_weight.to(device=self.device)) @@ -213,11 +217,15 @@ class disable_weight_init: error_msgs, ) + def forward_comfy_cast_weights(self, input): + weight, bias = cast_bias_weight(self, input) + return torch.nn.functional.linear(input, weight, bias) + def forward(self, input): + if self.comfy_cast_weights or len(self.weight_function) > 0 or len(self.bias_function) > 0: + return self.forward_comfy_cast_weights(input) weight = self.weight bias = self.bias - if self.comfy_cast_weights or len(self.weight_function) > 0 or len(self.bias_function) > 0: - weight, bias = cast_bias_weight(self, input) return torch.nn.functional.linear(input, weight, bias) class Conv1d(torch.nn.Conv1d, CastWeightBiasOp): @@ -391,7 +399,7 @@ try: except ImportError: pass -if CUBLAS_IS_AVAILABLE: # TODO check if this is actually faster I call BS +if CUBLAS_IS_AVAILABLE: class cublas_ops(disable_weight_init): class Linear(CublasLinear, disable_weight_init.Linear): def reset_parameters(self): @@ -408,6 +416,11 @@ op_class_list = [ if cls.__module__ == disable_weight_init.__module__ and name != "__class__" ] +op_fn_list = [ # conv_nd + fn for name, fn in inspect.getmembers(disable_weight_init, inspect.isroutine) + if not name.startswith('__') +] + def operator_factory(**factory_kwargs): class OpSet: pass @@ -417,9 +430,15 @@ def operator_factory(**factory_kwargs): for base_class in op_class_list: new_class = type(base_class.__name__, (base_class,), factory_kwargs) - setattr(op_set, base_class.__name__, new_class) + setattr(OpSet, base_class.__name__, new_class) - return op_set + for fn in op_fn_list: + setattr(OpSet, fn.__name__, fn) + + return OpSet + +# Backward Compatability (e.g. GGUF) +manual_cast = operator_factory(comfy_cast_weights=True) def pick_operations(weight_dtype=None, compute_dtype=None, load_device=None, fast_fp8=False, disable_fast_fp8=False): fp8_compute = (comfy.model_management.supports_fp8_compute(load_device) and not disable_fast_fp8) @@ -429,4 +448,5 @@ def pick_operations(weight_dtype=None, compute_dtype=None, load_device=None, fas weight_function = [] if weight_dtype is not None and compute_dtype is not None and not manual_cast: weight_function = [get_quantizer_with_constraints(weight_dtype)] - return operator_factory(comfy_cast_weights=manual_cast, use_dynamic_quantizer=use_dynamic_quantizer, fp8_compute=fp8_compute, weight_function=weight_function) + return operator_factory(comfy_cast_weights=manual_cast, use_dynamic_quantizer=use_dynamic_quantizer, + fp8_compute=fp8_compute, weight_function=weight_function)