From d8e8f47ce5c9aa4fe9fc338a25c9a011b7112769 Mon Sep 17 00:00:00 2001
From: Disty0 <disty@disty.xyz>
Date: Wed, 28 May 2025 15:18:39 +0300
Subject: [PATCH] SDNQ add an option to toggle quantize with GPU

---
 CHANGELOG.md                |  5 +++--
 modules/model_quant.py      |  2 +-
 modules/model_quant_sdnq.py | 17 +++++++++++++----
 modules/shared.py           |  1 +
 4 files changed, 18 insertions(+), 7 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index f9d3fa8ce..ea5672a57 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -13,14 +13,15 @@
   - Set the default quant mode to `pre`  
   - Use per token input quant with int8 matmul  
   - Implement better layer hijacks  
+  - Add an option to toggle quantize with GPU  
   - Fix Conv quant  
   - Fix lora weight change  
-  * Fix high RAM usage with pre mode  
+  - Fix high RAM usage with pre mode  
 - **IPEX**  
   - Disabe Dynamic Attention by default on PyTorch 2.7  
   - Remove GradScaler hijack and use torch.amp.GradScaler instead  
 - **Feature**  
-  - TeaCache support for HiDream I1 
+  - TeaCache support for HiDream I1  
 - **Changes**
   - Set the default attention optimizer to Scaled-Dot-Product on all backends  
   - Enable Dynamic attention for Scaled-Dot-Product with ROCm, DirectML, MPS and CPU backends  
diff --git a/modules/model_quant.py b/modules/model_quant.py
index b38451414..96579c4eb 100644
--- a/modules/model_quant.py
+++ b/modules/model_quant.py
@@ -349,7 +349,7 @@ def sdnq_quantize_model(model, op=None, sd_model=None, do_gc=True):
             quant_last_model_name = None
             quant_last_model_device = None
         model.to(devices.device)
-    elif shared.opts.diffusers_offload_mode == "model":
+    elif shared.opts.diffusers_offload_mode != "none":
         model = model.to(devices.cpu)
     if do_gc:
         devices.torch_gc(force=True)
diff --git a/modules/model_quant_sdnq.py b/modules/model_quant_sdnq.py
index ac7ce71f6..771029069 100644
--- a/modules/model_quant_sdnq.py
+++ b/modules/model_quant_sdnq.py
@@ -96,14 +96,19 @@ def sdnq_quantize_layer(layer, weights_dtype="int8", torch_dtype=None, group_siz
         if shared.opts.diffusers_offload_mode in {"none", "model"}:
             return_device = devices.device
         elif pre_mode:
-            if shared.opts.device_map != "gpu":
+            if shared.opts.sdnq_quantize_with_gpu:
                 return_device = devices.cpu
-            else:
+            elif shared.opts.device_map == "gpu":
                 return_device = devices.device
+            else:
+                return_device = layer.weight.device
         else:
             return_device = layer.weight.device
         if not pre_mode:
-            layer.weight.data = layer.weight.to(devices.device).to(dtype=torch.float32)
+            if shared.opts.sdnq_quantize_with_gpu:
+                layer.weight.data = layer.weight.to(devices.device).to(dtype=torch.float32)
+            else:
+                layer.weight.data = layer.weight.to(dtype=torch.float32)
 
         if dtype_dict[weights_dtype]["is_unsigned"]:
             scale, zero_point = get_scale_asymmetric(layer.weight, reduction_axes, weights_dtype)
@@ -493,7 +498,11 @@ class SDNQQuantizer(DiffusersQuantizer):
     ):
         # load the model params to target_device first
         layer, _ = get_module_from_name(model, param_name)
-        layer.weight = torch.nn.Parameter(param_value.to(devices.device).to(dtype=torch.float32), requires_grad=False)
+        if shared.opts.sdnq_quantize_with_gpu:
+            param_value = param_value.to(devices.device).to(dtype=torch.float32)
+        else:
+            param_value = param_value.to(target_device).to(dtype=torch.float32)
+        layer.weight = torch.nn.Parameter(param_value, requires_grad=False)
         layer = sdnq_quantize_layer(
             layer,
             weights_dtype=self.quantization_config.weights_dtype,
diff --git a/modules/shared.py b/modules/shared.py
index c329ef1fb..82e290902 100644
--- a/modules/shared.py
+++ b/modules/shared.py
@@ -524,6 +524,7 @@ options_templates.update(options_section(('quantization', "Quantization Settings
     "sdnq_decompress_fp32": OptionInfo(False, "Decompress using full precision", gr.Checkbox, {"visible": native}),
     "sdnq_decompress_compile": OptionInfo(devices.has_triton(), "Decompress using torch.compile", gr.Checkbox, {"visible": native}),
     "sdnq_use_quantized_matmul": OptionInfo(False, "Use quantized MatMul", gr.Checkbox, {"visible": native}),
+    "sdnq_quantize_with_gpu": OptionInfo(False, "Quantize with the GPU", gr.Checkbox, {"visible": native}),
     "sdnq_quantize_shuffle_weights": OptionInfo(False, "Shuffle weights in post mode", gr.Checkbox, {"visible": native}),
 
     "bnb_quantization_sep": OptionInfo("<h2>BitsAndBytes</h2>", "", gr.HTML),