Add sdnq_modules_dtype_dict

CHANGELOG.md

@@ -17,7 +17,7 @@ A quick service release with several important hotfixes, but also adding support
     default is `rust` as new `xet` is known to cause issues  
   - support for `flux.1-kontext` lora  
   - support for `qwen-image` lora  
-  - add `sdnq_modules_to_not_convert` option to ui settings  
+  - add `sdnq_modules_to_not_convert` and `sdnq_modules_dtype_dict` option to ui settings  
 - **UI**
   - new artwork for reference models in networks  
     thanks @liutyi  

modules/model_quant.py

@@ -2,6 +2,7 @@ import os
 import re
 import sys
 import copy
+import json
 import time
 import diffusers
 import transformers
@@ -142,6 +143,22 @@ def create_sdnq_config(kwargs = None, allow: bool = True, module: str = 'Model',
         if len(sdnq_modules_to_not_convert) > 0:
             modules_to_not_convert.extend(sdnq_modules_to_not_convert)
 
+        try:
+            if len(shared.opts.sdnq_modules_dtype_dict) > 2:
+                sdnq_modules_dtype_dict = shared.opts.sdnq_modules_dtype_dict
+                if "{" not in sdnq_modules_dtype_dict:
+                    sdnq_modules_dtype_dict = "{" + sdnq_modules_dtype_dict + "}"
+                sdnq_modules_dtype_dict = json.loads(bytes(sdnq_modules_dtype_dict, 'utf-8'))
+                for key, value in sdnq_modules_dtype_dict.items():
+                    if isinstance(value, str):
+                        value = [m.strip() for m in re.split(';|,| ', value) if len(m.strip()) > 1]
+                    if key not in modules_dtype_dict.keys():
+                        modules_dtype_dict[key] = value
+                    else:
+                        modules_dtype_dict[key].extend(value)
+        except Exception as e:
+            log.warning(f'Quantization: SDNQ failed to parse sdnq_modules_dtype_dict: {e}')
+
         quantization_device, return_device = get_sdnq_devices()
 
         sdnq_config = SDNQConfig(
@@ -155,7 +172,7 @@ def create_sdnq_config(kwargs = None, allow: bool = True, module: str = 'Model',
             quantization_device=quantization_device,
             return_device=return_device,
             modules_to_not_convert=modules_to_not_convert,
-            modules_dtype_dict=modules_dtype_dict,
+            modules_dtype_dict=modules_dtype_dict.copy(),
         )
         log.debug(f'Quantization: module="{module}" type=sdnq mode=pre dtype={weights_dtype} matmul={shared.opts.sdnq_use_quantized_matmul} group_size={shared.opts.sdnq_quantize_weights_group_size} quant_conv={shared.opts.sdnq_quantize_conv_layers} matmul_conv={shared.opts.sdnq_use_quantized_matmul_conv} dequantize_fp32={shared.opts.sdnq_dequantize_fp32} quantize_with_gpu={shared.opts.sdnq_quantize_with_gpu} quantization_device={quantization_device} return_device={return_device} device_map={shared.opts.device_map} offload_mode={shared.opts.diffusers_offload_mode} non_blocking={shared.opts.diffusers_offload_nonblocking} modules_to_not_convert={modules_to_not_convert} modules_dtype_dict={modules_dtype_dict}')
         if kwargs is None:
@@ -184,7 +201,7 @@ def check_nunchaku(module: str = ''):
     return True
 
 
-def create_config(kwargs = None, allow: bool = True, module: str = 'Model', modules_to_not_convert = [], modules_dtype_dict = {}):
+def create_config(kwargs = None, allow: bool = True, module: str = 'Model', modules_to_not_convert: list = [], modules_dtype_dict: dict = {}):
     if kwargs is None:
         kwargs = {}
     kwargs = create_sdnq_config(kwargs, allow=allow, module=module, modules_to_not_convert=modules_to_not_convert, modules_dtype_dict=modules_dtype_dict)
@@ -417,6 +434,22 @@ def sdnq_quantize_model(model, op=None, sd_model=None, do_gc: bool = True, weigh
     if len(sdnq_modules_to_not_convert) > 0:
         modules_to_not_convert.extend(sdnq_modules_to_not_convert)
 
+    try:
+        if len(shared.opts.sdnq_modules_dtype_dict) > 2:
+            sdnq_modules_dtype_dict = shared.opts.sdnq_modules_dtype_dict
+            if "{" not in sdnq_modules_dtype_dict:
+                sdnq_modules_dtype_dict = "{" + sdnq_modules_dtype_dict + "}"
+            sdnq_modules_dtype_dict = json.loads(bytes(sdnq_modules_dtype_dict, 'utf-8'))
+            for key, value in sdnq_modules_dtype_dict.items():
+                if isinstance(value, str):
+                    value = [m.strip() for m in re.split(';|,| ', value) if len(m.strip()) > 1]
+                if key not in modules_dtype_dict.keys():
+                    modules_dtype_dict[key] = value
+                else:
+                    modules_dtype_dict[key].extend(value)
+    except Exception as e:
+        log.warning(f'Quantization: SDNQ failed to parse sdnq_modules_dtype_dict: {e}')
+
     model.eval()
     backup_embeddings = None
     if hasattr(model, "get_input_embeddings"):
@@ -436,7 +469,7 @@ def sdnq_quantize_model(model, op=None, sd_model=None, do_gc: bool = True, weigh
         quantization_device=quantization_device,
         return_device=return_device,
         modules_to_not_convert=modules_to_not_convert,
-        modules_dtype_dict=modules_dtype_dict,
+        modules_dtype_dict=modules_dtype_dict.copy(),
         op=op,
     )
     t1 = time.time()

modules/sdnq/__init__.py

@@ -211,7 +211,6 @@ def apply_sdnq_to_module(model, weights_dtype="int8", torch_dtype=None, group_si
                                     weights_dtype = "u" + weights_dtype
                         else:
                             weights_dtype = key
-                        break
 
             module = sdnq_quantize_layer(
                 module,
@@ -326,7 +325,6 @@ class SDNQQuantizer(DiffusersQuantizer):
                                 weights_dtype = "u" + weights_dtype
                     else:
                         weights_dtype = key
-                    break
 
         if self.quantization_config.return_device is not None:
             return_device = self.quantization_config.return_device

modules/shared.py

@@ -183,6 +183,7 @@ options_templates.update(options_section(("quantization", "Model Quantization"),
     "sdnq_quantize_weights_mode": OptionInfo("int8", "Quantization type", gr.Dropdown, {"choices": sdnq_quant_modes}),
     "sdnq_quantize_weights_mode_te": OptionInfo("Same as model", "Quantization type for Text Encoders", gr.Dropdown, {"choices": ['Same as model'] + sdnq_quant_modes}),
     "sdnq_modules_to_not_convert": OptionInfo("", "Modules to not convert"),
+    "sdnq_modules_dtype_dict": OptionInfo("{}", "Modules dtype dict"),
     "sdnq_quantize_weights_group_size": OptionInfo(0, "Group size", gr.Slider, {"minimum": -1, "maximum": 4096, "step": 1}),
     "sdnq_quantize_conv_layers": OptionInfo(False, "Quantize convolutional layers", gr.Checkbox),
     "sdnq_dequantize_compile": OptionInfo(devices.has_triton(), "Dequantize using torch.compile", gr.Checkbox),