From 48b5d56ba4ffb5b0160a2cee0e7378b37ea5e1a2 Mon Sep 17 00:00:00 2001
From: Disty0 <disty@disty.xyz>
Date: Wed, 26 Nov 2025 21:08:15 +0300
Subject: [PATCH] Enable or disable quantized matmul on pre-quant models

---
 modules/sd_models.py   | 11 ++++++++++-
 modules/sdnq/loader.py |  6 +++---
 2 files changed, 13 insertions(+), 4 deletions(-)

diff --git a/modules/sd_models.py b/modules/sd_models.py
index d37f7b154..7a49993fb 100644
--- a/modules/sd_models.py
+++ b/modules/sd_models.py
@@ -14,7 +14,7 @@ from modules import timer, paths, shared, shared_items, modelloader, devices, sc
 from modules.memstats import memory_stats
 from modules.modeldata import model_data
 from modules.sd_checkpoint import CheckpointInfo, select_checkpoint, list_models, checkpoints_list, checkpoint_titles, get_closest_checkpoint_match, model_hash, update_model_hashes, setup_model, write_metadata, read_metadata_from_safetensors # pylint: disable=unused-import
-from modules.sd_offload import disable_offload, set_diffuser_offload, apply_balanced_offload, set_accelerate # pylint: disable=unused-import
+from modules.sd_offload import get_module_names, disable_offload, set_diffuser_offload, apply_balanced_offload, set_accelerate # pylint: disable=unused-import
 from modules.sd_models_utils import NoWatermark, get_signature, get_call, path_to_repo, patch_diffuser_config, convert_to_faketensors, read_state_dict, get_state_dict_from_checkpoint, apply_function_to_model # pylint: disable=unused-import
 
 
@@ -164,6 +164,15 @@ def set_diffuser_options(sd_model, vae=None, op:str='model', offload:bool=True,
         shared.log.quiet(quiet, f'Setting {op}: channels-last=True')
         sd_model.unet.to(memory_format=torch.channels_last)
 
+    for module_name in get_module_names(sd_model):
+        module = getattr(sd_model, module_name, None)
+        if hasattr(module, "quantization_config") and getattr(module.quantization_config, "quant_method", None) == "sdnq":
+            if module.quantization_config.use_quantized_matmul != shared.opts.sdnq_use_quantized_matmul:
+                from sdnq.loader import apply_sdnq_options_to_model
+                shared.log.debug(f'Setting {op} {module_name}: sdnq_use_quantized_matmul={shared.opts.sdnq_use_quantized_matmul}')
+                module = apply_sdnq_options_to_model(module, use_quantized_matmul=shared.opts.sdnq_use_quantized_matmul)
+                setattr(sd_model, module_name, module)
+
     if offload:
         set_diffuser_offload(sd_model, op, quiet)
 
diff --git a/modules/sdnq/loader.py b/modules/sdnq/loader.py
index d1d1e0f47..f79edfb90 100644
--- a/modules/sdnq/loader.py
+++ b/modules/sdnq/loader.py
@@ -134,7 +134,7 @@ def load_sdnq_model(model_path: str, model_cls: ModelMixin = None, file_name: st
 
     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)
+        model = apply_sdnq_options_to_model(model, dtype=dtype, dequantize_fp32=dequantize_fp32, use_quantized_matmul=use_quantized_matmul)
     return model
 
 
@@ -154,7 +154,7 @@ def post_process_model(model):
     return model
 
 
-def apply_options_to_model(model, dtype: torch.dtype = None, dequantize_fp32: bool = None, use_quantized_matmul: bool = None):
+def apply_sdnq_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:
         if dtype is not None and getattr(model, "dtype", torch.float32) != torch.float32:
@@ -199,5 +199,5 @@ def apply_options_to_model(model, dtype: torch.dtype = None, dequantize_fp32: bo
                 module.forward = module.forward.__get__(module, module.__class__)
             setattr(model, module_name, module)
         else:
-            setattr(model, module_name, apply_options_to_model(module, dtype=dtype, dequantize_fp32=dequantize_fp32, use_quantized_matmul=use_quantized_matmul))
+            setattr(model, module_name, apply_sdnq_options_to_model(module, dtype=dtype, dequantize_fp32=dequantize_fp32, use_quantized_matmul=use_quantized_matmul))
     return model