Make SDNQ not depended on quantization_config.json and fix invalid quantization_config getting attached to the model on load

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__:

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,

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

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

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:')