From 93f28f07acda29012c77091b1f9e9f6ea0aaa1fb Mon Sep 17 00:00:00 2001 From: Disty0 Date: Fri, 7 Nov 2025 18:11:21 +0300 Subject: [PATCH] Make SDNQ not depended on quantization_config.json and fix invalid quantization_config getting attached to the model on load --- modules/sd_detect.py | 11 ++++++ modules/sd_models.py | 9 ++++- modules/sdnq/loader.py | 59 ++++++++++++++-------------- modules/sdnq/quantizer.py | 81 ++++++++++++++++++++++++++++----------- pipelines/generic.py | 14 +++++-- 5 files changed, 118 insertions(+), 56 deletions(-) diff --git a/modules/sd_detect.py b/modules/sd_detect.py index 2e169ae11..f95d9b9d6 100644 --- a/modules/sd_detect.py +++ b/modules/sd_detect.py @@ -156,11 +156,22 @@ def guess_by_diffusers(fn, current_guess): if folder.endswith('quantization_config.json'): is_quant = True break + if folder.endswith('config.json'): + quantization_config = shared.readfile(folder, silent=True).get("quantization_config", None) + if quantization_config is not None: + is_quant = True + break if os.path.isdir(folder): for f in os.listdir(folder): + f = os.path.join(folder, f) if f.endswith('quantization_config.json'): is_quant = True break + if f.endswith('config.json'): + quantization_config = shared.readfile(f, silent=True).get("quantization_config", None) + if quantization_config is not None: + is_quant = True + break pipelines = shared_items.get_pipelines() for k, v in pipelines.items(): if v is not None and v.__name__ == pipeline.__name__: diff --git a/modules/sd_models.py b/modules/sd_models.py index bd52deada..b40cf4aea 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -559,11 +559,16 @@ def load_diffuser_file(model_type, pipeline, checkpoint_info, diffusers_load_con def load_sdnq_module(fn: str, module_name: str, load_method: str): from modules import sdnq t0 = time.time() + quantization_config = None quantization_config_path = os.path.join(fn, module_name, 'quantization_config.json') - if not os.path.exists(quantization_config_path): + model_config_path = os.path.join(fn, module_name, 'config.json') + if os.path.exists(quantization_config_path): + quantization_config = shared.readfile(quantization_config_path, silent=True) + elif os.path.exists(model_config_path): + quantization_config = shared.readfile(model_config_path, silent=True).get("quantization_config", None) + if quantization_config is None: return None, module_name, 0 model_name = os.path.join(fn, module_name) - quantization_config = shared.readfile(quantization_config_path, silent=True) try: module = sdnq.load_sdnq_model( model_path=model_name, diff --git a/modules/sdnq/loader.py b/modules/sdnq/loader.py index 749af9fb4..cc3b35bb2 100644 --- a/modules/sdnq/loader.py +++ b/modules/sdnq/loader.py @@ -2,8 +2,9 @@ import os import json import torch from diffusers.models.modeling_utils import ModelMixin -from .common import dtype_dict, use_tensorwise_fp8_matmul, use_contiguous_mm -from .quantizer import SDNQConfig, sdnq_post_load_quant + +from .common import dtype_dict, use_tensorwise_fp8_matmul +from .quantizer import SDNQConfig, sdnq_post_load_quant, prepare_weight_for_matmul, prepare_svd_for_matmul from .dequantizer import dequantize_symmetric, re_quantize_int8, re_quantize_fp8 from .forward import get_forward_func from .file_loader import load_files @@ -68,11 +69,16 @@ def load_sdnq_model(model_path: str, model_cls: ModelMixin = None, file_name: st with init_empty_weights(): if quantization_config is None: - try: - with open(os.path.join(model_path, "quantization_config.json"), "r", encoding="utf-8") as f: + quantization_config_path = os.path.join(model_path, "quantization_config.json") + model_config_path = os.path.join(model_path, "config.json") + if os.path.exists(quantization_config_path): + with open(quantization_config_path, "r", encoding="utf-8") as f: quantization_config = json.load(f) - except Exception: - quantization_config = {} + elif os.path.exists(model_config_path): + with open(model_config_path, "r", encoding="utf-8") as f: + quantization_config = json.load(f).get("quantization_config", None) + if quantization_config is None: + raise ValueError(f"Cannot determine quantization_config for {model_path}, please provide quantization_config argument") if model_config is None: try: @@ -127,11 +133,26 @@ def load_sdnq_model(model_path: str, model_cls: ModelMixin = None, file_name: st model.load_state_dict(state_dict, assign=True) del state_dict + model = post_process_model(model) if (dtype is not None) or (dequantize_fp32 is not None) or (use_quantized_matmul is not None): model = apply_options_to_model(model, dtype=dtype, dequantize_fp32=dequantize_fp32, use_quantized_matmul=use_quantized_matmul) return model +def post_process_model(model): + has_children = list(model.children()) + if not has_children: + return model + for module in model.children(): + if hasattr(module, "sdnq_dequantizer"): + if module.sdnq_dequantizer.use_quantized_matmul and not module.sdnq_dequantizer.re_quantize_for_matmul: + module.weight.data = prepare_weight_for_matmul(module.weight) + if module.svd_up is not None: + module.svd_up.data, module.svd_down.data = prepare_svd_for_matmul(module.svd_up, module.svd_down, module.sdnq_dequantizer.use_quantized_matmul) + module = post_process_model(module) + return model + + def apply_options_to_model(model, dtype: torch.dtype = None, dequantize_fp32: bool = None, use_quantized_matmul: bool = None): has_children = list(model.children()) if not has_children: @@ -168,30 +189,12 @@ def apply_options_to_model(model, dtype: torch.dtype = None, dequantize_fp32: bo if use_tensorwise_fp8_matmul: module.scale.data = module.scale.to(dtype=scale_dtype) elif not module.sdnq_dequantizer.re_quantize_for_matmul: - module.weight.data, module.scale.data = module.weight.t_(), module.scale.t_() + module.scale.t_() + module.weight.t_() if use_quantized_matmul: - if use_contiguous_mm: - module.weight.data = module.weight.contiguous() - elif module.weight.is_contiguous(): - module.weight.data = module.weight.t_().contiguous().t_() + module.weight.data = prepare_weight_for_matmul(module.weight) if module.svd_up is not None: - module.svd_up.data = module.svd_up.t_() - module.svd_down.data = module.svd_down.t_() - if use_quantized_matmul: - if use_contiguous_mm: - module.svd_up.data = module.svd_up.contiguous() - module.svd_down.data = module.svd_down.contiguous() - else: - if module.svd_up.is_contiguous(): - module.svd_up.data = module.svd_up.t_().contiguous().t_() - if module.svd_up.is_contiguous(): - module.svd_down.data = module.svd_down.t_().contiguous().t_() - else: - module.svd_up.data = module.svd_up.contiguous() - if use_contiguous_mm: - module.svd_down.data = module.svd_down.contiguous() - elif module.svd_down.is_contiguous(): - module.svd_down.data = module.svd_down.t_().contiguous().t_() + module.svd_up.data, module.svd_down.data = prepare_svd_for_matmul(module.svd_up.t_(), module.svd_down.t_(), use_quantized_matmul) module.sdnq_dequantizer.use_quantized_matmul = use_quantized_matmul module.forward = get_forward_func(module.__class__.__name__, use_quantized_matmul, dtype_dict[module.sdnq_dequantizer.weights_dtype]["is_integer"], use_tensorwise_fp8_matmul) module.forward = module.forward.__get__(module, module.__class__) diff --git a/modules/sdnq/quantizer.py b/modules/sdnq/quantizer.py index 14ecfcc92..b0bae850e 100644 --- a/modules/sdnq/quantizer.py +++ b/modules/sdnq/quantizer.py @@ -68,6 +68,25 @@ def apply_svdquant(weight: torch.FloatTensor, rank: int = 32, niter: int = 8) -> return weight, svd_up, svd_down +def prepare_weight_for_matmul(weight: torch.Tensor) -> torch.Tensor: + if use_contiguous_mm: + weight = weight.contiguous() + elif weight.is_contiguous(): + weight = weight.t_().contiguous().t_() + return weight + + +def prepare_svd_for_matmul(svd_up: torch.FloatTensor, svd_down: torch.FloatTensor, use_quantized_matmul: bool) -> Tuple[torch.FloatTensor, torch.FloatTensor]: + if svd_up is not None: + if use_quantized_matmul: + svd_up = prepare_weight_for_matmul(svd_up) + else: + svd_up = svd_up.contiguous() + if svd_down is not None: + svd_down = prepare_weight_for_matmul(svd_down) + return svd_up, svd_down + + def check_param_name_in(param_name: str, param_list: List[str]) -> bool: split_param_name = param_name.split(".") for param in param_list: @@ -212,20 +231,7 @@ def sdnq_quantize_layer(layer, weights_dtype="int8", torch_dtype=None, group_siz if use_quantized_matmul: svd_up = svd_up.t_() svd_down = svd_down.t_() - if use_contiguous_mm: - svd_up = svd_up.contiguous() - svd_down = svd_down.contiguous() - else: - if svd_up.is_contiguous(): - svd_up = svd_up.t_().contiguous().t_() - if svd_down.is_contiguous(): - svd_down = svd_down.t_().contiguous().t_() - else: - svd_up = svd_up.contiguous() - if use_contiguous_mm: - svd_down = svd_down.contiguous() - elif svd_down.is_contiguous(): - svd_down = svd_down.t_().contiguous().t_() + svd_up, svd_down = prepare_svd_for_matmul(svd_up, svd_down, use_quantized_matmul) except Exception: svd_up, svd_down = None, None else: @@ -295,10 +301,7 @@ def sdnq_quantize_layer(layer, weights_dtype="int8", torch_dtype=None, group_siz if use_quantized_matmul and not re_quantize_for_matmul: scale.t_() layer.weight.t_() - if use_contiguous_mm: - layer.weight.data = layer.weight.contiguous() - elif layer.weight.is_contiguous(): - layer.weight.data = layer.weight.t_().contiguous().t_() + layer.weight.data = prepare_weight_for_matmul(layer.weight) if not use_tensorwise_fp8_matmul and not dtype_dict[weights_dtype]["is_integer"]: scale = scale.to(dtype=torch.float32) @@ -418,6 +421,13 @@ def sdnq_post_load_quant( modules_dtype_dict: Dict[str, List[str]] = None, op=None, ): + if modules_to_not_convert is None: + modules_to_not_convert = [] + if modules_dtype_dict is None: + modules_dtype_dict = {} + + modules_to_not_convert = modules_to_not_convert.copy() + modules_dtype_dict = modules_dtype_dict.copy() if add_skip_keys: model, modules_to_not_convert, modules_dtype_dict = add_module_skip_keys(model, modules_to_not_convert, modules_dtype_dict) @@ -438,7 +448,7 @@ def sdnq_post_load_quant( quantization_device=quantization_device, return_device=return_device, modules_to_not_convert=modules_to_not_convert, - modules_dtype_dict=modules_dtype_dict.copy(), + modules_dtype_dict=modules_dtype_dict, op=op, ) model.quantization_config = SDNQConfig( @@ -455,12 +465,15 @@ def sdnq_post_load_quant( quantization_device=quantization_device, return_device=return_device, modules_to_not_convert=modules_to_not_convert, - modules_dtype_dict=modules_dtype_dict.copy(), + modules_dtype_dict=modules_dtype_dict, ) if hasattr(model, "config"): try: model.config.quantization_config = model.quantization_config + except Exception: + pass + try: model.config["quantization_config"] = model.quantization_config.to_dict() except Exception: pass @@ -543,6 +556,14 @@ class SDNQQuantizer(DiffusersQuantizer, HfQuantizer): param_value = param_value.clone() else: param_value = param_value.to(target_device, dtype=return_dtype) + + if tensor_name == "weight" and layer.sdnq_dequantizer.use_quantized_matmul and not layer.sdnq_dequantizer.re_quantize_for_matmul: + param_value = prepare_weight_for_matmul(param_value) + elif tensor_name == "svd_up": + param_value, _ = prepare_svd_for_matmul(param_value, None, layer.sdnq_dequantizer.use_quantized_matmul) + elif tensor_name == "svd_down": + _, param_value = prepare_svd_for_matmul(None, param_value, layer.sdnq_dequantizer.use_quantized_matmul) + param_value = torch.nn.Parameter(param_value, requires_grad=False) setattr(layer, tensor_name, param_value) return @@ -626,6 +647,9 @@ class SDNQQuantizer(DiffusersQuantizer, HfQuantizer): if hasattr(model, "config"): try: model.config.quantization_config = self.quantization_config + except Exception: + pass + try: model.config["quantization_config"] = self.quantization_config.to_dict() except Exception: pass @@ -655,8 +679,17 @@ class SDNQQuantizer(DiffusersQuantizer, HfQuantizer): del model.quantization_method if hasattr(model, "quantization_config"): del model.quantization_config - if hasattr(model, "config") and hasattr(model.config, "quantization_config"): - del model.config.quantization_config + if hasattr(model, "config"): + try: + if hasattr(model.config, "quantization_config"): + del model.config.quantization_config + except Exception: + pass + try: + if hasattr(model.config, "pop"): + model.config.pop("quantization_config", None) + except Exception: + pass return model def is_serializable(self, *args, **kwargs) -> bool: # pylint: disable=unused-argument, invalid-overridden-method @@ -772,6 +805,7 @@ class SDNQConfig(QuantizationConfigMixin): 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: + self.modules_dtype_dict = self.modules_dtype_dict.copy() for key, value in self.modules_dtype_dict.items(): if isinstance(value, str): value = [value] @@ -782,6 +816,9 @@ class SDNQConfig(QuantizationConfigMixin): 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)}") + self.modules_to_not_convert = self.modules_to_not_convert.copy() + self.modules_dtype_dict = self.modules_dtype_dict.copy() + def to_dict(self): dct = self.__dict__.copy() # make serializable dct["quantization_device"] = str(dct["quantization_device"]) if dct["quantization_device"] is not None else None diff --git a/pipelines/generic.py b/pipelines/generic.py index 896d3c5c2..3ae0a154d 100644 --- a/pipelines/generic.py +++ b/pipelines/generic.py @@ -76,8 +76,11 @@ def load_transformer(repo_id, cls_name, load_config=None, subfolder="transformer if shared.opts.diffusers_offload_mode != 'none' and transformer is not None: sd_models.move_model(transformer, devices.cpu) - if (transformer is not None) and (quant_type is not None) and (quant_args.get('quantization_config', None) is not None): # attach quantization_config - transformer.quantization_config = quant_args.get('quantization_config', None) + if transformer is not None and not hasattr(transformer, 'quantization_config'): # attach quantization_config + if hasattr(transformer, 'config') and hasattr(transformer.config, 'quantization_config'): + transformer.quantization_config = transformer.config.quantization_config + elif (quant_type is not None) and (quant_args.get('quantization_config', None) is not None): + transformer.quantization_config = quant_args.get('quantization_config', None) except Exception as e: shared.log.error(f'Load model: transformer="{repo_id}" cls={cls_name.__name__} {e}') errors.display(e, 'Load:') @@ -209,8 +212,11 @@ def load_text_encoder(repo_id, cls_name, load_config=None, subfolder="text_encod if shared.opts.diffusers_offload_mode != 'none' and text_encoder is not None: sd_models.move_model(text_encoder, devices.cpu) - if (text_encoder is not None) and (quant_type is not None) and (quant_args.get('quantization_config', None) is not None): # attach quantization_config - text_encoder.quantization_config = quant_args.get('quantization_config', None) + if text_encoder is not None and not hasattr(text_encoder, 'quantization_config'): # attach quantization_config + if hasattr(text_encoder, 'config') and hasattr(text_encoder.config, 'quantization_config'): + text_encoder.quantization_config = text_encoder.config.quantization_config + elif (quant_type is not None) and (quant_args.get('quantization_config', None) is not None): + text_encoder.quantization_config = quant_args.get('quantization_config', None) except Exception as e: shared.log.error(f'Load model: text_encoder="{repo_id}" cls={cls_name.__name__} {e}') errors.display(e, 'Load:')