From 7aa1bfdc7001c758a5fd144fc891a0c076b54010 Mon Sep 17 00:00:00 2001
From: Disty0 <disty@disty.xyz>
Date: Tue, 2 Dec 2025 01:01:51 +0300
Subject: [PATCH] Add get_modules_to_not_convert from transformers v5

---
 modules/sdnq/common.py    |  4 ++--
 modules/sdnq/quantizer.py | 12 ++++++++++++
 2 files changed, 14 insertions(+), 2 deletions(-)

diff --git a/modules/sdnq/common.py b/modules/sdnq/common.py
index e8d3c9ae4..bdb50a03c 100644
--- a/modules/sdnq/common.py
+++ b/modules/sdnq/common.py
@@ -5,7 +5,7 @@ import torch
 
 from modules import shared, devices
 
-sdnq_version = "0.1.1"
+sdnq_version = "0.1.2"
 
 dtype_dict = {
     "int32": {"min": -2147483648, "max": 2147483647, "num_bits": 32, "sign": 1, "exponent": 0, "mantissa": 31, "target_dtype": torch.int32, "torch_dtype": torch.int32, "storage_dtype": torch.int32, "is_unsigned": False, "is_integer": True, "is_packed": False},
@@ -17,7 +17,7 @@ dtype_dict = {
     "int4": {"min": -8, "max": 7, "num_bits": 4, "sign": 1, "exponent": 0, "mantissa": 3, "target_dtype": "int4", "torch_dtype": torch.int8, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": True, "is_packed": True},
     "int3": {"min": -4, "max": 3, "num_bits": 3, "sign": 1, "exponent": 0, "mantissa": 2, "target_dtype": "int3", "torch_dtype": torch.int8, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": True, "is_packed": True},
     "int2": {"min": -2, "max": 1, "num_bits": 2, "sign": 1, "exponent": 0, "mantissa": 1, "target_dtype": "int2", "torch_dtype": torch.int8, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": True, "is_packed": True},
-    "uint32": {"min": 0, "max": 4294967295, "num_bits": 32, "sign": 0, "exponent": 0, "mantissa": 31, "target_dtype": torch.uint32, "torch_dtype": torch.uint32, "storage_dtype": torch.uint32, "is_unsigned": True, "is_integer": True, "is_packed": False},
+    "uint32": {"min": 0, "max": 4294967295, "num_bits": 32, "sign": 0, "exponent": 0, "mantissa": 32, "target_dtype": torch.uint32, "torch_dtype": torch.uint32, "storage_dtype": torch.uint32, "is_unsigned": True, "is_integer": True, "is_packed": False},
     "uint16": {"min": 0, "max": 65535, "num_bits": 16, "sign": 0, "exponent": 0, "mantissa": 16, "target_dtype": torch.uint16, "torch_dtype": torch.uint16, "storage_dtype": torch.uint16, "is_unsigned": True, "is_integer": True, "is_packed": False},
     "uint8": {"min": 0, "max": 255, "num_bits": 8, "sign": 0, "exponent": 0, "mantissa": 8, "target_dtype": torch.uint8, "torch_dtype": torch.uint8, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": True, "is_packed": False},
     "uint7": {"min": 0, "max": 127, "num_bits": 7, "sign": 0, "exponent": 0, "mantissa": 7, "target_dtype": "uint7", "torch_dtype": torch.uint8, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": True, "is_packed": True},
diff --git a/modules/sdnq/quantizer.py b/modules/sdnq/quantizer.py
index 954ed106d..7320865dd 100644
--- a/modules/sdnq/quantizer.py
+++ b/modules/sdnq/quantizer.py
@@ -727,6 +727,12 @@ class SDNQQuantizer(DiffusersQuantizer, HfQuantizer):
         self.torch_dtype = torch_dtype
         return torch_dtype
 
+    def update_dtype(self, dtype: torch.dtype = None) -> torch.dtype:
+        """
+        needed for transformers compatibilty, returns self.update_torch_dtype
+        """
+        return self.update_torch_dtype(dtype)
+
     def _process_model_before_weight_loading( # pylint: disable=arguments-differ
         self,
         model,
@@ -746,6 +752,8 @@ class SDNQQuantizer(DiffusersQuantizer, HfQuantizer):
         if self.quantization_config.add_skip_keys:
             if keep_in_fp32_modules is not None:
                 self.quantization_config.modules_to_not_convert.extend(keep_in_fp32_modules)
+            if hasattr(self, "get_modules_to_not_convert") and hasattr(model, "tie_weights"):
+                self.quantization_config.modules_to_not_convert.extend(self.get_modules_to_not_convert(model, add_default_skips=True))
             model, self.quantization_config.modules_to_not_convert, self.quantization_config.modules_dtype_dict = add_module_skip_keys(
                 model, self.quantization_config.modules_to_not_convert, self.quantization_config.modules_dtype_dict
             )
@@ -799,6 +807,10 @@ class SDNQQuantizer(DiffusersQuantizer, HfQuantizer):
     def is_trainable(self):
         return self.quantization_config.is_training
 
+    @property
+    def is_qat_trainable(self) -> bool:
+        return self.is_trainable()
+
     @property
     def is_compileable(self):
         return True