From e49814098ecb0b78ed4f1c0baebaf740dd2e19a8 Mon Sep 17 00:00:00 2001 From: Disty0 Date: Wed, 20 Aug 2025 14:58:54 +0300 Subject: [PATCH] Add sdnq_modules_dtype_dict --- CHANGELOG.md | 2 +- modules/model_quant.py | 39 ++++++++++++++++++++++++++++++++++++--- modules/sdnq/__init__.py | 2 -- modules/shared.py | 1 + 4 files changed, 38 insertions(+), 6 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 5f57556bb..b76091cae 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/modules/model_quant.py b/modules/model_quant.py index 8b5991e8c..3275a6e91 100644 --- a/modules/model_quant.py +++ b/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() diff --git a/modules/sdnq/__init__.py b/modules/sdnq/__init__.py index 8f260db78..642506842 100644 --- a/modules/sdnq/__init__.py +++ b/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 diff --git a/modules/shared.py b/modules/shared.py index 0cffb049a..ad53b2a51 100644 --- a/modules/shared.py +++ b/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),