diff --git a/modules/sdnq/common.py b/modules/sdnq/common.py index 8d0114e3e..57a4f0aaf 100644 --- a/modules/sdnq/common.py +++ b/modules/sdnq/common.py @@ -35,6 +35,7 @@ linear_types = {"Linear"} conv_types = {"Conv1d", "Conv2d", "Conv3d"} conv_transpose_types = {"ConvTranspose1d", "ConvTranspose2d", "ConvTranspose3d"} allowed_types = set.union(linear_types, conv_types, conv_transpose_types) +accepted_weights = set(dtype_dict.keys()) use_torch_compile = shared.opts.sdnq_dequantize_compile # this setting requires a full restart of the webui to apply is_rdna2 = bool(devices.backend == "rocm" and int(getattr(torch.cuda.get_device_properties(devices.device), "gcnArchName", "gfx0000")[3:]) < 1100) diff --git a/modules/sdnq/quantizer.py b/modules/sdnq/quantizer.py index e56ecb43c..56c6c6d7f 100644 --- a/modules/sdnq/quantizer.py +++ b/modules/sdnq/quantizer.py @@ -12,7 +12,7 @@ from diffusers.quantizers.quantization_config import QuantizationConfigMixin from diffusers.utils import get_module_from_name from modules import devices, shared -from .common import dtype_dict, use_tensorwise_fp8_matmul, allowed_types, conv_types, conv_transpose_types, use_contiguous_mm +from .common import dtype_dict, accepted_weights, use_tensorwise_fp8_matmul, allowed_types, conv_types, conv_transpose_types, use_contiguous_mm from .dequantizer import dequantizer_dict, dequantize_sdnq_model from .forward import get_forward_func @@ -411,7 +411,7 @@ def sdnq_post_load_quant( class SDNQQuantizer(DiffusersQuantizer, HfQuantizer): r""" - Diffusers Quantizer for SDNQ + Diffusers and Transformers Quantizer for SDNQ """ requires_parameters_quantization = True @@ -751,7 +751,7 @@ class SDNQConfig(QuantizationConfigMixin): svd_rank (`int`, *optional*, defaults to `32`): The rank size used for the SVDQuant algorithm. use_svd (`bool`, *optional*, defaults to `False`): - Enabling this option will use SVDQuant algorithm. + Enabling this option will use SVDQuant algorithm on top of SDNQ quantization. quant_conv (`bool`, *optional*, defaults to `False`): Enabling this option will quantize the convolutional layers in UNet models too. use_quantized_matmul (`bool`, *optional*, defaults to `False`): @@ -815,17 +815,32 @@ class SDNQConfig(QuantizationConfigMixin): r""" Safety checker that arguments are correct """ - accepted_weights = ["int8", "int7", "int6", "int5", "int4", "int3", "int2", "uint8", "uint7", "uint6", "uint5", "uint4", "uint3", "uint2", "uint1", "bool", "float8_e4m3fn", "float8_e4m3fnuz", "float8_e5m2", "float8_e5m2fnuz"] if self.weights_dtype not in accepted_weights: - raise ValueError(f"Only support weights in {accepted_weights} but found {self.weights_dtype}") + raise ValueError(f"SDNQ only support weights in {accepted_weights} but found {self.weights_dtype}") if self.modules_to_not_convert is None: self.modules_to_not_convert = [] - elif not isinstance(self.modules_to_not_convert, list): + elif isinstance(self.modules_to_not_convert, str): self.modules_to_not_convert = [self.modules_to_not_convert] + elif isinstance(self.modules_to_not_convert, tuple): + self.modules_to_not_convert = list(self.modules_to_not_convert) + elif not isinstance(self.modules_to_not_convert, list): + raise ValueError(f"modules_to_not_convert must be a list but got {type(self.modules_to_not_convert)}") if self.modules_dtype_dict is None: self.modules_dtype_dict = {} + elif not isinstance(self.modules_dtype_dict, dict): + raise ValueError(f"modules_dtype_dict must be a dict but got {type(self.modules_dtype_dict)}") + elif len(self.modules_dtype_dict.keys()) > 0: + for key, value in self.modules_dtype_dict.items(): + if isinstance(value, str): + value = [value] + self.modules_dtype_dict[key] = value + elif isinstance(value, tuple): + value = list(value) + self.modules_dtype_dict[key] = value + if not isinstance(key, str) or not isinstance(value, list): + raise ValueError(f"modules_dtype_dict must be a dictionary of strings and lists but got {type(key)} and {type(value)}") def to_dict(self): dct = self.__dict__.copy() # make serializable