Merge pull request #3982 from vladmandic/dev

Dev

CHANGELOG.md

@@ -1,5 +1,46 @@
 # Change Log for SD.Next
 
+## Update for 2025-06-15
+
+- **Feature**  
+  - Support for Python 3.13  
+
+- **Changes**  
+  - Increase the medvram mode threshold from 8GB to 12GB  
+  - Set CPU backend to use FP32 by default  
+  - Relax Python version checks for Zluda  
+  - Make VAE options not require model reload  
+  - Add warning about incompatible attention processors  
+
+- **Torch**  
+  - Set default to `torch==2.7.1`  
+  - Force upgrade pip when installing Torch  
+
+- **ROCm**  
+  - Support ROCm 6.4 with `--use-nightly`  
+  - Don't override user set gfx version  
+  - Don't override gfx version with RX 9000  
+  - Fix flash-atten repo  
+
+- **SDNQ Quantization**  
+  - Add group size support for convolutional layers  
+  - Add quantized matmul support for for convolutional layers  
+  - Add 7-bit, 5-bit and 3-bit quantization support  
+  - Add separate quant mode option for Text Encoders  
+  - Fix forced FP32 with tensorwise FP8 matmul  
+  - Fix PyTorch <= 2.4 compatibility with FP8 matmul  
+  - Fix VAE with conv quant  
+  - Don't ignore the Quantize with GPU option with offload mode `none` and `model`  
+
+- **Fixes**  
+  - Meissonic with multiple generators  
+  - OmniGen with new transformers  
+  - Invalid attention processors  
+  - PixArt Sigma Small and Large loading  
+  - TAESD previews with PixArt and Lumina 2  
+  - VAE Tiling with non-default tile sizes  
+  - Lumina 2 with IPEX  
+
 ## Update for 2025-06-02
 
 ### Highlights for 2025-06-02
@@ -31,7 +72,6 @@ Take a look at [Docs](https://github.com/vladmandic/sdnext/wiki/Docs), [Hints](h
     - `INT4` -> `uint4`  
   - Add `float8_e4m3fn`, `float8_e5m2`, `float8_e4m3fnuz`, `float8_e5m2fnuz`, `int6`, `uint6`, `int2`, `uint2` and `uint1` support  
   - Add quantized matmul support for `float8_e4m3fn` and `float8_e5m2`  
-  - Add group size support for convolutional layers  
   - Set the default quant mode to `pre`  
   - Use per token input quant with int8 and fp8 quantized matmul  
   - Implement better layer hijacks  

installer.py

@@ -513,7 +513,7 @@ def get_platform():
 def check_python(supported_minors=[], experimental_minors=[], reason=None):
     if supported_minors is None or len(supported_minors) == 0:
         supported_minors = [9, 10, 11, 12]
-        experimental_minors = []
+        experimental_minors = [13]
     t_start = time.time()
     if args.quick:
         return
@@ -546,7 +546,7 @@ def check_diffusers():
     t_start = time.time()
     if args.skip_all or args.skip_git or args.experimental:
         return
-    sha = '6508da6f06a0da1054ae6a808d0025c04b70f0e8' # diffusers commit hash
+    sha = '8adc6003ba4dbf5b61bb4f1ce571e9e55e145a99' # diffusers commit hash
     pkg = pkg_resources.working_set.by_key.get('diffusers', None)
     minor = int(pkg.version.split('.')[1] if pkg is not None else 0)
     cur = opts.get('diffusers_version', '') if minor > 0 else ''
@@ -582,7 +582,7 @@ def install_cuda():
         cmd = os.environ.get('TORCH_COMMAND', '--upgrade --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cu128 --extra-index-url https://download.pytorch.org/whl/nightly/cu126')
     else:
         # cmd = os.environ.get('TORCH_COMMAND', 'torch==2.6.0+cu126 torchvision==0.21.0+cu126 --index-url https://download.pytorch.org/whl/cu126')
-        cmd = os.environ.get('TORCH_COMMAND', 'torch==2.7.0+cu128 torchvision==0.22.0+cu128 --index-url https://download.pytorch.org/whl/cu128')
+        cmd = os.environ.get('TORCH_COMMAND', 'torch==2.7.1+cu128 torchvision==0.22.1+cu128 --index-url https://download.pytorch.org/whl/cu128')
     return cmd
 
 
@@ -633,7 +633,7 @@ def install_rocm_zluda():
     log.info(msg)
 
     if sys.platform == "win32": # TODO install: enable ROCm for windows when available
-        check_python(supported_minors=[10, 11], reason='ZLUDA backend requires Python 3.10 or 3.11')
+        #check_python(supported_minors=[9, 10, 11, 12, 13], reason='ZLUDA backend requires a Python version between 3.9 and 3.13')
 
         if args.device_id is not None:
             if os.environ.get('HIP_VISIBLE_DEVICES', None) is not None:
@@ -655,7 +655,7 @@ def install_rocm_zluda():
         if error is None:
             try:
                 zluda_installer.load()
-                torch_command = os.environ.get('TORCH_COMMAND', 'torch==2.7.0 torchvision --index-url https://download.pytorch.org/whl/cu118')
+                torch_command = os.environ.get('TORCH_COMMAND', 'torch==2.7.1+cu118 torchvision==0.22.1+cu118 --index-url https://download.pytorch.org/whl/cu118')
             except Exception as e:
                 error = e
                 log.warning(f'Failed to load ZLUDA: {e}')
@@ -663,22 +663,24 @@ def install_rocm_zluda():
             log.info('Using CPU-only torch')
             torch_command = os.environ.get('TORCH_COMMAND', 'torch torchvision')
     else:
-        check_python(supported_minors=[9, 10, 11, 12], reason='ROCm backend requires a Python version between 3.9 and 3.12')
+        #check_python(supported_minors=[9, 10, 11, 12, 13], reason='ROCm backend requires a Python version between 3.9 and 3.13')
 
         if os.environ.get("TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL", None) is None:
             os.environ.setdefault('TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL', '1')
 
         if args.use_nightly:
-            if rocm.version is None or float(rocm.version) >= 6.3: # assume the latest if version check fails
+            if rocm.version is None or float(rocm.version) >= 6.4: # assume the latest if version check fails
+                torch_command = os.environ.get('TORCH_COMMAND', '--upgrade --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/rocm6.4')
+            elif rocm.version == "6.3":
                 torch_command = os.environ.get('TORCH_COMMAND', '--upgrade --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/rocm6.3')
             else: # oldest rocm version on nightly is 6.2.4
                 torch_command = os.environ.get('TORCH_COMMAND', '--upgrade --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/rocm6.2.4')
         else:
             if rocm.version is None or float(rocm.version) >= 6.3: # assume the latest if version check fails
-                torch_command = os.environ.get('TORCH_COMMAND', 'torch==2.7.0+rocm6.3 torchvision==0.22.0+rocm6.3 --index-url https://download.pytorch.org/whl/rocm6.3')
+                torch_command = os.environ.get('TORCH_COMMAND', 'torch==2.7.1+rocm6.3 torchvision==0.22.1+rocm6.3 --index-url https://download.pytorch.org/whl/rocm6.3')
             elif rocm.version == "6.2":
-                # use rocm 6.2.4 instead of 6.2 as torch==2.7.0+rocm6.2 doesn't exists
-                torch_command = os.environ.get('TORCH_COMMAND', 'torch==2.7.0+rocm6.2.4 torchvision==0.22.0+rocm6.2.4 --index-url https://download.pytorch.org/whl/rocm6.2.4')
+                # use rocm 6.2.4 instead of 6.2 as torch==2.7.1+rocm6.2 doesn't exists
+                torch_command = os.environ.get('TORCH_COMMAND', 'torch==2.7.1+rocm6.2.4 torchvision==0.22.1+rocm6.2.4 --index-url https://download.pytorch.org/whl/rocm6.2.4')
             elif rocm.version == "6.1":
                 torch_command = os.environ.get('TORCH_COMMAND', 'torch==2.6.0+rocm6.1 torchvision==0.21.0+rocm6.1 --index-url https://download.pytorch.org/whl/rocm6.1')
             elif rocm.version == "6.0":
@@ -696,22 +698,21 @@ def install_rocm_zluda():
             log.debug(f'ROCm hipBLASLt: arch={device.name} available={device.blaslt_supported}')
             rocm.set_blaslt_enabled(device.blaslt_supported)
 
-    if device is None:
-        log.debug('ROCm: HSA_OVERRIDE_GFX_VERSION auto config skipped')
+    if device is None or os.environ.get("HSA_OVERRIDE_GFX_VERSION", None) is not None:
+        log.info(f'ROCm: HSA_OVERRIDE_GFX_VERSION auto config skipped: device={device.name if device is not None else None} version={os.environ.get("HSA_OVERRIDE_GFX_VERSION", None)}')
     else:
         gfx_ver = device.get_gfx_version()
         if gfx_ver is not None:
             os.environ.setdefault('HSA_OVERRIDE_GFX_VERSION', gfx_ver)
-        else:
-            log.warning(f'ROCm: device={device.name} could not auto-detect HSA version')
+            log.info(f'ROCm: HSA_OVERRIDE_GFX_VERSION config overridden: device={device.name} version={os.environ.get("HSA_OVERRIDE_GFX_VERSION", None)}')
 
     ts('amd', t_start)
     return torch_command
 
 
-def install_ipex(torch_command):
+def install_ipex():
     t_start = time.time()
-    check_python(supported_minors=[9, 10, 11, 12], reason='IPEX backend requires a Python version between 3.9 and 3.12')
+    #check_python(supported_minors=[9, 10, 11, 12, 13], reason='IPEX backend requires a Python version between 3.9 and 3.13')
     args.use_ipex = True # pylint: disable=attribute-defined-outside-init
     log.info('IPEX: Intel OneAPI toolkit detected')
 
@@ -736,20 +737,20 @@ def install_ipex(torch_command):
     if args.use_nightly:
         torch_command = os.environ.get('TORCH_COMMAND', '--upgrade --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/xpu')
     else:
-        torch_command = os.environ.get('TORCH_COMMAND', 'torch==2.7.0+xpu torchvision==0.22.0+xpu --index-url https://download.pytorch.org/whl/xpu')
+        torch_command = os.environ.get('TORCH_COMMAND', 'torch==2.7.1+xpu torchvision==0.22.1+xpu --index-url https://download.pytorch.org/whl/xpu')
 
     ts('ipex', t_start)
     return torch_command
 
 
-def install_openvino(torch_command):
+def install_openvino():
     t_start = time.time()
-    check_python(supported_minors=[9, 10, 11, 12], reason='OpenVINO backend requires a Python version between 3.9 and 3.12')
+    #check_python(supported_minors=[9, 10, 11, 12, 13], reason='OpenVINO backend requires a Python version between 3.9 and 3.13')
     log.info('OpenVINO: selected')
     if sys.platform == 'darwin':
-        torch_command = os.environ.get('TORCH_COMMAND', 'torch==2.7.0 torchvision==0.22.0')
+        torch_command = os.environ.get('TORCH_COMMAND', 'torch==2.7.1 torchvision==0.22.1')
     else:
-        torch_command = os.environ.get('TORCH_COMMAND', 'torch==2.7.0+cpu torchvision==0.22.0+cpu --index-url https://download.pytorch.org/whl/cpu')
+        torch_command = os.environ.get('TORCH_COMMAND', 'torch==2.7.1+cpu torchvision==0.22.1+cpu --index-url https://download.pytorch.org/whl/cpu')
 
     install(os.environ.get('OPENVINO_COMMAND', 'openvino==2025.1.0'), 'openvino')
     install(os.environ.get('NNCF_COMMAND', 'nncf==2.16.0'), 'nncf')
@@ -840,15 +841,15 @@ def check_torch():
         elif is_rocm_available and (args.use_rocm or args.use_zluda): # prioritize rocm
             torch_command = install_rocm_zluda()
         elif allow_ipex and args.use_ipex: # prioritize ipex
-            torch_command = install_ipex(torch_command)
+            torch_command = install_ipex()
         elif allow_openvino and args.use_openvino: # prioritize openvino
-            torch_command = install_openvino(torch_command)
+            torch_command = install_openvino()
         elif is_cuda_available:
             torch_command = install_cuda()
         elif is_rocm_available:
             torch_command = install_rocm_zluda()
         elif is_ipex_available:
-            torch_command = install_ipex(torch_command)
+            torch_command = install_ipex()
         else:
             machine = platform.machine()
             if sys.platform == 'darwin':
@@ -867,6 +868,7 @@ def check_torch():
     if 'torch' in torch_command and not args.version:
         if not installed('torch'):
             log.info(f'Torch: download and install in progress... cmd="{torch_command}"')
+            install('--upgrade pip', 'pip', reinstall=True) # pytorch rocm is too large for older pip
         install(torch_command, 'torch torchvision', quiet=True)
     else:
         try:
@@ -1155,8 +1157,8 @@ def ensure_base_requirements():
 def install_optional():
     t_start = time.time()
     log.info('Installing optional requirements...')
-    install('basicsr')
-    install('gfpgan')
+    install('git+https://github.com/Disty0/BasicSR@2b6a12c28e0c81bfb13b7e984144f0b0f5461484', 'basicsr')
+    install('git+https://github.com/Disty0/GFPGAN@09b1190eabbc77e5f15c61fa7c38a2064b403e20', 'gfpgan')
     install('clean-fid')
     install('pillow-jxl-plugin==1.3.3', ignore=True)
     install('optimum-quanto==0.2.7', ignore=True)
@@ -1188,6 +1190,16 @@ def install_requirements():
         pr.enable()
     if args.skip_requirements and not args.requirements:
         return
+    if int(sys.version_info.minor) >= 13:
+        install('audioop-lts')
+        # gcc 15 patch
+        backup_cmake_policy = os.environ.get('CMAKE_POLICY_VERSION_MINIMUM', None)
+        backup_cxxflags = os.environ.get('CXXFLAGS', None)
+        os.environ.setdefault('CMAKE_POLICY_VERSION_MINIMUM', '3.5')
+        os.environ.setdefault('CXXFLAGS', '-include cstdint')
+        install('git+https://github.com/google/sentencepiece#subdirectory=python', 'sentencepiece')
+        os.environ.setdefault('CMAKE_POLICY_VERSION_MINIMUM', backup_cmake_policy)
+        os.environ.setdefault('CXXFLAGS', backup_cxxflags)
     if not installed('diffusers', quiet=True): # diffusers are not installed, so run initial installation
         global quick_allowed # pylint: disable=global-statement
         quick_allowed = False

modules/devices.py

@@ -331,7 +331,7 @@ def test_fp16():
     if fp16_ok is not None:
         return fp16_ok
     if opts.cuda_dtype != 'FP16': # don't override if the user sets it
-        if sys.platform == "darwin" or backend == 'openvino': # override
+        if sys.platform == "darwin" or backend in {'openvino', 'cpu'}: # override
             fp16_ok = False
             return fp16_ok
         elif backend == 'rocm':
@@ -362,7 +362,7 @@ def test_bf16():
     if bf16_ok is not None:
         return bf16_ok
     if opts.cuda_dtype != 'BF16': # don't override if the user sets it
-        if sys.platform == "darwin" or backend == 'openvino' or backend == 'directml': # override
+        if sys.platform == "darwin" or backend in {'openvino', 'directml', 'cpu'}: # override
             bf16_ok = False
             return bf16_ok
         elif backend == 'rocm' or backend == 'zluda':
@@ -426,8 +426,6 @@ def override_ipex_math():
 
 def set_sdpa_params():
     try:
-        if opts.cross_attention_optimization != "Scaled-Dot-Product":
-            return
         try:
             global sdpa_original # pylint: disable=global-statement
             if sdpa_original is not None:

modules/intel/ipex/__init__.py

@@ -39,7 +39,6 @@ def ipex_init(): # pylint: disable=too-many-statements
             torch.cuda.is_available = torch.xpu.is_available
             torch.cuda.is_initialized = torch.xpu.is_initialized
             torch.cuda.is_current_stream_capturing = lambda: False
-            torch.cuda.set_device = torch.xpu.set_device
             torch.cuda.stream = torch.xpu.stream
             torch.cuda.Event = torch.xpu.Event
             torch.cuda.Stream = torch.xpu.Stream

modules/intel/ipex/diffusers.py

@@ -81,14 +81,46 @@ def get_1d_sincos_pos_embed_from_grid(embed_dim, pos, output_type="np"):
     return emb
 
 
+def apply_rotary_emb(x, freqs_cis, use_real: bool = True, use_real_unbind_dim: int = -1):
+    if use_real:
+        cos, sin = freqs_cis  # [S, D]
+        cos = cos[None, None]
+        sin = sin[None, None]
+        cos, sin = cos.to(x.device), sin.to(x.device)
+
+        if use_real_unbind_dim == -1:
+            # Used for flux, cogvideox, hunyuan-dit
+            x_real, x_imag = x.reshape(*x.shape[:-1], -1, 2).unbind(-1)  # [B, S, H, D//2]
+            x_rotated = torch.stack([-x_imag, x_real], dim=-1).flatten(3)
+        elif use_real_unbind_dim == -2:
+            # Used for Stable Audio, OmniGen, CogView4 and Cosmos
+            x_real, x_imag = x.reshape(*x.shape[:-1], 2, -1).unbind(-2)  # [B, S, H, D//2]
+            x_rotated = torch.cat([-x_imag, x_real], dim=-1)
+        else:
+            raise ValueError(f"`use_real_unbind_dim={use_real_unbind_dim}` but should be -1 or -2.")
+
+        out = (x.float() * cos + x_rotated.float() * sin).to(x.dtype)
+        return out
+    else:
+        # used for lumina
+        # force cpu with Alchemist
+        x_rotated = torch.view_as_complex(x.to("cpu").float().reshape(*x.shape[:-1], -1, 2))
+        freqs_cis = freqs_cis.to("cpu").unsqueeze(2)
+        x_out = torch.view_as_real(x_rotated * freqs_cis).flatten(3)
+        return x_out.type_as(x).to(x.device)
+
+
 def ipex_diffusers(device_supports_fp64=False):
     diffusers.utils.torch_utils.fourier_filter = fourier_filter
     if not device_supports_fp64:
         # get around lazy imports
+        from diffusers.models import embeddings as diffusers_embeddings # pylint: disable=import-error, unused-import # noqa: F401
         from diffusers.models import transformers as diffusers_transformers # pylint: disable=import-error, unused-import # noqa: F401
         from diffusers.models import controlnets as diffusers_controlnets # pylint: disable=import-error, unused-import # noqa: F401
         diffusers.models.embeddings.get_1d_sincos_pos_embed_from_grid = get_1d_sincos_pos_embed_from_grid
         diffusers.models.embeddings.FluxPosEmbed = FluxPosEmbed
+        diffusers.models.embeddings.apply_rotary_emb = apply_rotary_emb
         diffusers.models.transformers.transformer_flux.FluxPosEmbed = FluxPosEmbed
+        diffusers.models.transformers.transformer_lumina2.apply_rotary_emb = apply_rotary_emb
         diffusers.models.controlnets.controlnet_flux.FluxPosEmbed = FluxPosEmbed
         diffusers.models.transformers.transformer_hidream_image.rope = hidream_rope

modules/intel/ipex/hijacks.py

@@ -379,6 +379,12 @@ def torch_cuda_device(device):
     else:
         return torch.xpu.device(device)
 
+@wraps(torch.cuda.set_device)
+def torch_cuda_set_device(device):
+    if check_cuda(device):
+        torch.xpu.set_device(return_xpu(device))
+    else:
+        torch.xpu.set_device(device)
 
 # torch.Generator has to be a class for isinstance checks
 original_torch_Generator = torch.Generator
@@ -412,6 +418,7 @@ def ipex_hijacks():
     torch.load = torch_load
     torch.cuda.synchronize = torch_cuda_synchronize
     torch.cuda.device = torch_cuda_device
+    torch.cuda.set_device = torch_cuda_set_device
 
     torch.Generator = torch_Generator
     torch._C.Generator = torch_Generator

modules/lora/lora_apply.py

@@ -45,8 +45,8 @@ def network_backup_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.n
                     self.network_weights_backup = True
                 else:
                     self.network_weights_backup = weight.clone().to(devices.cpu)
-                    if hasattr(self, "sdnq_decompressor"):
-                        self.sdnq_decompressor_backup = self.sdnq_decompressor.to(devices.cpu)
+                    if hasattr(self, "sdnq_dequantizer"):
+                        self.sdnq_dequantizer_backup = self.sdnq_dequantizer.to(devices.cpu)
 
         if bias_backup is None:
             if getattr(self, 'bias', None) is not None:
@@ -79,8 +79,8 @@ def network_calc_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn.
             continue
         try:
             t0 = time.time()
-            if hasattr(self, "sdnq_decompressor"):
-                weight = self.sdnq_decompressor.to(devices.device)(self.weight.to(devices.device), skip_quantized_matmul=self.sdnq_decompressor.use_quantized_matmul)
+            if hasattr(self, "sdnq_dequantizer"):
+                weight = self.sdnq_dequantizer.to(devices.device)(self.weight.to(devices.device), skip_quantized_matmul=self.sdnq_dequantizer.use_quantized_matmul)
             else:
                 weight = self.weight.to(devices.device) # must perform calc on gpu due to performance
             updown, ex_bias = module.calc_updown(weight)
@@ -136,18 +136,30 @@ def network_add_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn.G
             # weight._quantize(devices.device) / weight.to(device=device)
         except Exception as e:
             shared.log.error(f'Network load: type=LoRA quant=bnb cls={self.__class__.__name__} type={self.quant_type} blocksize={self.blocksize} state={vars(self.quant_state)} weight={self.weight} bias={lora_weights} {e}')
-    elif not bias and hasattr(self, "sdnq_decompressor"):
+    elif not bias and hasattr(self, "sdnq_dequantizer"):
         try:
-            from modules.model_quant_sdnq import sdnq_quantize_layer
-            if hasattr(self, "sdnq_decompressor_backup"):
-                sdnq_decompressor = self.sdnq_decompressor_backup.to(devices.device)
+            from modules.sdnq import sdnq_quantize_layer
+            if hasattr(self, "sdnq_dequantizer_backup"):
+                sdnq_dequantizer = self.sdnq_dequantizer_backup.to(devices.device)
             else:
-                sdnq_decompressor = self.sdnq_decompressor.to(devices.device)
-            dequant_weight = sdnq_decompressor(model_weights.to(devices.device), skip_quantized_matmul=sdnq_decompressor.use_quantized_matmul)
+                sdnq_dequantizer = self.sdnq_dequantizer.to(devices.device)
+            dequant_weight = sdnq_dequantizer(model_weights.to(devices.device), skip_quantized_matmul=sdnq_dequantizer.use_quantized_matmul)
             new_weight = dequant_weight.to(devices.device, dtype=torch.float32) + lora_weights.to(devices.device, dtype=torch.float32)
             self.weight = torch.nn.Parameter(new_weight, requires_grad=False)
-            self.sdnq_decompressor = None
-            self = sdnq_quantize_layer(self, sdnq_decompressor.weights_dtype, torch_dtype=devices.dtype, group_size=shared.opts.sdnq_quantize_weights_group_size, quant_conv=shared.opts.sdnq_quantize_conv_layers, use_quantized_matmul=shared.opts.sdnq_use_quantized_matmul, param_name=getattr(self, 'network_layer_name', None))
+            self.sdnq_dequantizer = None
+            self = sdnq_quantize_layer(
+                self,
+                sdnq_dequantizer.weights_dtype,
+                torch_dtype=devices.dtype,
+                group_size=shared.opts.sdnq_quantize_weights_group_size,
+                quant_conv=shared.opts.sdnq_quantize_conv_layers,
+                use_quantized_matmul=shared.opts.sdnq_use_quantized_matmul,
+                use_quantized_matmul_conv=shared.opts.sdnq_use_quantized_matmul_conv,
+                dequantize_fp32=shared.opts.sdnq_dequantize_fp32,
+                quantization_device=devices.device,
+                return_device=device,
+                param_name=getattr(self, 'network_layer_name', None),
+            )
             self = self.to(device)
             weight = None
             del dequant_weight
@@ -214,8 +226,8 @@ def network_apply_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn
             network_add_weights(self, model_weights=weights_backup, lora_weights=updown, deactivate=deactivate, device=device, bias=False)
         else:
             self.weight = torch.nn.Parameter(weights_backup.to(device), requires_grad=False)
-            if hasattr(self, "sdnq_decompressor_backup"):
-                self.sdnq_decompressor = self.sdnq_decompressor_backup.to(device)
+            if hasattr(self, "sdnq_dequantizer_backup"):
+                self.sdnq_dequantizer = self.sdnq_dequantizer_backup.to(device)
 
     if bias_backup is not None:
         self.bias = None

modules/meissonic/scheduler.py

@@ -23,8 +23,12 @@ from diffusers.schedulers.scheduling_utils import SchedulerMixin
 
 
 def gumbel_noise(t, generator=None):
-    device = generator.device if generator is not None else t.device
-    noise = torch.zeros_like(t, device=device).uniform_(0, 1, generator=generator).to(t.device)
+    noise = []
+    noise_shape = t.shape[1:]
+    for i in range(len(generator)):
+        device = generator[i].device if generator[i] is not None else t.device
+        noise.append(torch.zeros(noise_shape, device=device, dtype=t.dtype).uniform_(0, 1, generator=generator[i]).to(t.device))
+    noise = torch.stack(noise, dim=0)
     return -torch.log((-torch.log(noise.clamp(1e-20))).clamp(1e-20))
 
 
@@ -100,14 +104,20 @@ class Scheduler(SchedulerMixin, ConfigMixin):
         unknown_map = sample == self.config.mask_token_id
 
         probs = model_output.softmax(dim=-1)
-
         device = probs.device
-        probs_ = probs.to(generator.device) if generator is not None else probs  # handles when generator is on CPU
-        if probs_.device.type == "cpu" and probs_.dtype != torch.float32:
-            probs_ = probs_.float()  # multinomial is not implemented for cpu half precision
-        probs_ = probs_.reshape(-1, probs.size(-1))
-        pred_original_sample = torch.multinomial(probs_, 1, generator=generator).to(device=device)
-        pred_original_sample = pred_original_sample[:, 0].view(*probs.shape[:-1])
+        probs_view_shape = probs.shape[1:-1]
+        if not isinstance(generator, list):
+            generator = [generator] * probs.size(0)
+        elif isinstance(generator, list) and len(generator) == 1 and len(generator) != probs.size(0):
+            generator = generator * probs.size(0)
+
+        pred_original_sample = []
+        for i in range(len(generator)):
+            probs_ = probs[i].to(generator[i].device) if generator[i] is not None else probs[i] # handles when generator is on CPU
+            if probs_.device.type == "cpu" and probs_.dtype != torch.float32:
+                probs_ = probs_.float()  # multinomial is not implemented for cpu half precision
+            pred_original_sample.append(torch.multinomial(probs_, 1, generator=generator[i]).to(device=device).view(*probs_view_shape))
+        pred_original_sample = torch.stack(pred_original_sample, dim=0)
         pred_original_sample = torch.where(unknown_map, pred_original_sample, sample)
 
         if timestep == 0:
@@ -163,7 +173,7 @@ class Scheduler(SchedulerMixin, ConfigMixin):
 
         mask_indices = (
             torch.rand(
-                sample.shape, device=generator.device if generator is not None else sample.device, generator=generator
+                sample.shape, device=generator[0].device if generator[0] is not None else sample.device, generator=generator
             ).to(sample.device)
             < mask_ratio
         )

modules/model_flux.py

@@ -14,7 +14,6 @@ debug = shared.log.trace if os.environ.get('SD_LOAD_DEBUG', None) is not None el
 def load_flux_quanto(checkpoint_info):
     transformer, text_encoder_2 = None, None
     quanto = model_quant.load_quanto('Load model: type=FLUX')
-    quanto.tensor.qbits.QBitsTensor.create = lambda *args, **kwargs: quanto.tensor.qbits.QBitsTensor(*args, **kwargs)
 
     if isinstance(checkpoint_info, str):
         repo_path = checkpoint_info
@@ -36,11 +35,12 @@ def load_flux_quanto(checkpoint_info):
         quanto.requantize(transformer, state_dict, quantization_map, device=torch.device("cpu"))
         if shared.opts.diffusers_eval:
             transformer.eval()
-        if transformer.dtype != devices.dtype:
+        transformer_dtype = transformer.dtype
+        if transformer_dtype != devices.dtype:
             try:
                 transformer = transformer.to(dtype=devices.dtype)
             except Exception:
-                shared.log.error(f"Load model: type=FLUX Failed to cast transformer to {devices.dtype}, set dtype to {transformer.dtype}")
+                shared.log.error(f"Load model: type=FLUX Failed to cast transformer to {devices.dtype}, set dtype to {transformer_dtype}")
     except Exception as e:
         shared.log.error(f"Load model: type=FLUX failed to load Quanto transformer: {e}")
         if debug:
@@ -63,11 +63,12 @@ def load_flux_quanto(checkpoint_info):
         quanto.requantize(text_encoder_2, state_dict, quantization_map, device=torch.device("cpu"))
         if shared.opts.diffusers_eval:
             text_encoder_2.eval()
-        if text_encoder_2.dtype != devices.dtype:
+        text_encoder_2_dtype = text_encoder_2.dtype
+        if text_encoder_2_dtype != devices.dtype:
             try:
                 text_encoder_2 = text_encoder_2.to(dtype=devices.dtype)
             except Exception:
-                shared.log.error(f"Load model: type=FLUX Failed to cast text encoder to {devices.dtype}, set dtype to {text_encoder_2.dtype}")
+                shared.log.error(f"Load model: type=FLUX Failed to cast text encoder to {devices.dtype}, set dtype to {text_encoder_2_dtype}")
     except Exception as e:
         shared.log.error(f"Load model: type=FLUX failed to load Quanto text encoder: {e}")
         if debug:

modules/model_lumina.py

@@ -23,7 +23,7 @@ def load_lumina2(checkpoint_info, diffusers_load_config={}):
     transformer = diffusers.Lumina2Transformer2DModel.from_pretrained(
         repo_id,
         subfolder="transformer",
-        cache_dir=shared.opts.hfcache_dir,
+        cache_dir=shared.opts.diffusers_dir,
         **load_config,
         **quant_config,
     )
@@ -32,14 +32,13 @@ def load_lumina2(checkpoint_info, diffusers_load_config={}):
     text_encoder = transformers.AutoModel.from_pretrained(
         repo_id,
         subfolder="text_encoder",
-        cache_dir=shared.opts.hfcache_dir,
-        torch_dtype=devices.dtype,
+        cache_dir=shared.opts.diffusers_dir,
         **load_config,
         **quant_config,
     )
 
     load_config, quant_config = model_quant.get_dit_args(diffusers_load_config, allow_quant=False)
-    pipe = diffusers.Lumina2Text2ImgPipeline.from_pretrained(
+    pipe = diffusers.Lumina2Pipeline.from_pretrained(
         repo_id,
         cache_dir=shared.opts.diffusers_dir,
         text_encoder=text_encoder,

modules/model_pixart.py

@@ -1,11 +1,19 @@
 import transformers
 import diffusers
+from huggingface_hub import file_exists
 
 
 def load_pixart(checkpoint_info, diffusers_load_config={}):
     from modules import shared, devices, modelloader, sd_models, model_quant
     modelloader.hf_login()
     repo_id = sd_models.path_to_repo(checkpoint_info.name)
+    repo_id_tenc = repo_id
+    repo_id_pipe = repo_id
+
+    if not file_exists(repo_id_tenc, "text_encoder/config.json"):
+        repo_id_tenc = "PixArt-alpha/PixArt-Sigma-XL-2-1024-MS"
+    if not file_exists(repo_id_pipe, "model_index.json"):
+        repo_id_pipe = "PixArt-alpha/PixArt-Sigma-XL-2-1024-MS"
 
     load_args, quant_args = model_quant.get_dit_args(diffusers_load_config, module='Transformer')
     transformer = diffusers.PixArtTransformer2DModel.from_pretrained(
@@ -17,7 +25,7 @@ def load_pixart(checkpoint_info, diffusers_load_config={}):
     )
     load_args, quant_args = model_quant.get_dit_args(diffusers_load_config, module='TE', device_map=True)
     text_encoder = transformers.T5EncoderModel.from_pretrained(
-        repo_id,
+        repo_id_tenc,
         subfolder="text_encoder",
         cache_dir=shared.opts.hfcache_dir,
         **load_args,
@@ -26,7 +34,7 @@ def load_pixart(checkpoint_info, diffusers_load_config={}):
 
     load_args, _quant_args = model_quant.get_dit_args(diffusers_load_config, allow_quant=False)
     pipe = diffusers.PixArtSigmaPipeline.from_pretrained(
-        'PixArt-alpha/PixArt-Sigma-XL-2-1024-MS',
+        repo_id_pipe,
         cache_dir=shared.opts.diffusers_dir,
         transformer=transformer,
         text_encoder=text_encoder,

modules/model_quant.py

@@ -104,22 +104,45 @@ def create_quanto_config(kwargs = None, allow_quanto: bool = True, module: str =
 
 
 def create_sdnq_config(kwargs = None, allow_sdnq: bool = True, module: str = 'Model', weights_dtype: str = None):
-    from modules import shared
+    from modules import devices, shared
     if len(shared.opts.sdnq_quantize_weights) > 0 and (shared.opts.sdnq_quantize_mode == 'pre') and allow_sdnq:
         if 'Model' in shared.opts.sdnq_quantize_weights or (module is not None and module in shared.opts.sdnq_quantize_weights) or module == 'any':
-            from modules.model_quant_sdnq import SDNQQuantizer, SDNQConfig
+            from modules.sdnq import SDNQQuantizer, SDNQConfig
             diffusers.quantizers.auto.AUTO_QUANTIZER_MAPPING["sdnq"] = SDNQQuantizer
             transformers.quantizers.auto.AUTO_QUANTIZER_MAPPING["sdnq"] = SDNQQuantizer
             diffusers.quantizers.auto.AUTO_QUANTIZATION_CONFIG_MAPPING["sdnq"] = SDNQConfig
             transformers.quantizers.auto.AUTO_QUANTIZATION_CONFIG_MAPPING["sdnq"] = SDNQConfig
 
+            if weights_dtype is None:
+                if shared.opts.sdnq_quantize_weights_mode_te != "default" and module in {"TE", "LLM"}:
+                    weights_dtype = shared.opts.sdnq_quantize_weights_mode_te
+                else:
+                    weights_dtype = shared.opts.sdnq_quantize_weights_mode
+
+            if shared.opts.device_map == "gpu":
+                quantization_device = devices.device
+                return_device = devices.device
+            elif shared.opts.diffusers_offload_mode in {"none", "model"}:
+                quantization_device = devices.device if shared.opts.sdnq_quantize_with_gpu else devices.cpu
+                return_device = devices.device
+            elif shared.opts.sdnq_quantize_with_gpu:
+                quantization_device = devices.device
+                return_device = devices.cpu
+            else:
+                quantization_device = None
+                return_device = None
+
             sdnq_config = SDNQConfig(
-                weights_dtype=weights_dtype if weights_dtype is not None else shared.opts.sdnq_quantize_weights_mode,
+                weights_dtype=weights_dtype,
                 group_size=shared.opts.sdnq_quantize_weights_group_size,
                 quant_conv=shared.opts.sdnq_quantize_conv_layers,
                 use_quantized_matmul=shared.opts.sdnq_use_quantized_matmul,
+                use_quantized_matmul_conv=shared.opts.sdnq_use_quantized_matmul_conv,
+                dequantize_fp32=shared.opts.sdnq_dequantize_fp32,
+                quantization_device=quantization_device,
+                return_device=return_device,
             )
-            log.debug(f'Quantization: module="{module}" type=sdnq dtype={shared.opts.sdnq_quantize_weights_mode}')
+            log.debug(f'Quantization: module="{module}" type=sdnq dtype={weights_dtype}')
             if kwargs is None:
                 return sdnq_config
             else:
@@ -302,13 +325,13 @@ def apply_layerwise(sd_model, quiet:bool=False):
 def sdnq_quantize_model(model, op=None, sd_model=None, do_gc=True):
     global quant_last_model_name, quant_last_model_device # pylint: disable=global-statement
     from modules import devices, shared
-    from modules.model_quant_sdnq import apply_sdnq_to_module
+    from modules.sdnq import apply_sdnq_to_module
 
     model.eval()
 
     if model.__class__.__name__ in {"T5EncoderModel", "UMT5EncoderModel"}:
         import torch
-        from modules.model_quant_sdnq import SDNQ_T5DenseGatedActDense # T5DenseGatedActDense uses fp32
+        from modules.sdnq import SDNQ_T5DenseGatedActDense # T5DenseGatedActDense uses fp32
         for i in range(len(model.encoder.block)):
             model.encoder.block[i].layer[1].DenseReluDense = SDNQ_T5DenseGatedActDense(
                 model.encoder.block[i].layer[1].DenseReluDense,
@@ -319,13 +342,32 @@ def sdnq_quantize_model(model, op=None, sd_model=None, do_gc=True):
     if hasattr(model, "get_input_embeddings"):
         backup_embeddings = copy.deepcopy(model.get_input_embeddings())
 
+    if shared.opts.sdnq_quantize_weights_mode_te != "default" and op is not None and "text_encoder" in op:
+        weights_dtype = shared.opts.sdnq_quantize_weights_mode_te
+    else:
+        weights_dtype = shared.opts.sdnq_quantize_weights_mode
+
+    if shared.opts.diffusers_offload_mode in {"none", "model"}:
+        quantization_device = devices.device if shared.opts.sdnq_quantize_with_gpu else devices.cpu
+        return_device = devices.device
+    elif shared.opts.sdnq_quantize_with_gpu:
+        quantization_device = devices.device
+        return_device = getattr(model, "device", devices.cpu)
+    else:
+        quantization_device = None
+        return_device = None
+
     model = apply_sdnq_to_module(
         model,
-        weights_dtype=shared.opts.sdnq_quantize_weights_mode,
+        weights_dtype=weights_dtype,
         torch_dtype=devices.dtype,
         group_size=shared.opts.sdnq_quantize_weights_group_size,
         quant_conv=shared.opts.sdnq_quantize_conv_layers,
         use_quantized_matmul=shared.opts.sdnq_use_quantized_matmul,
+        use_quantized_matmul_conv=shared.opts.sdnq_use_quantized_matmul_conv,
+        dequantize_fp32=shared.opts.sdnq_dequantize_fp32,
+        quantization_device=quantization_device,
+        return_device=return_device,
         param_name=op,
     )
     model.quantization_method = 'SDNQ'

modules/model_quant_sdnq.py

@@ -1,853 +0,0 @@
-# pylint: disable=redefined-builtin,no-member,protected-access
-
-from typing import Any, Dict, List, Tuple, Optional, Union
-from dataclasses import dataclass
-from enum import Enum
-import sys
-import torch
-from diffusers.quantizers.base import DiffusersQuantizer
-from diffusers.quantizers.quantization_config import QuantizationConfigMixin
-from diffusers.utils import get_module_from_name
-from accelerate.utils import CustomDtype
-from modules import devices, shared
-
-torch_version = float(torch.__version__[:3])
-
-dtype_dict = {
-    "int8": {"min": -128, "max": 127, "num_bits": 8, "target_dtype": torch.int8, "torch_dtype": torch.int8, "storage_dtype": torch.int8, "is_unsigned": False, "is_integer": True},
-    "uint8": {"min": 0, "max": 255, "num_bits": 8, "target_dtype": torch.uint8, "torch_dtype": torch.uint8, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": True},
-    "int6": {"min": -32, "max": 31, "num_bits": 6, "target_dtype": torch.int8, "torch_dtype": torch.int8, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": True},
-    "uint6": {"min": 0, "max": 63, "num_bits": 6, "target_dtype": torch.uint8, "torch_dtype": torch.uint8, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": True},
-    "int4": {"min": -8, "max": 7, "num_bits": 4, "target_dtype": CustomDtype.INT4, "torch_dtype": torch.int8, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": True},
-    "uint4": {"min": 0, "max": 15, "num_bits": 4, "target_dtype": CustomDtype.INT4, "torch_dtype": torch.uint8, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": True},
-    "int2": {"min": -2, "max": 1, "num_bits": 2, "target_dtype": CustomDtype.INT2, "torch_dtype": torch.int8, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": True},
-    "uint2": {"min": 0, "max": 3, "num_bits": 2, "target_dtype": CustomDtype.INT2, "torch_dtype": torch.uint8, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": True},
-    "uint1": {"min": 0, "max": 1, "num_bits": 1, "target_dtype": torch.bool, "torch_dtype": torch.bool, "storage_dtype": torch.bool, "is_unsigned": True, "is_integer": True},
-    "float8_e4m3fn": {"min": -448, "max": 448, "num_bits": 8, "target_dtype": torch.float8_e4m3fn, "torch_dtype": torch.float8_e4m3fn, "storage_dtype": torch.float8_e4m3fn, "is_unsigned": False, "is_integer": False},
-    "float8_e5m2": {"min": -57344, "max": 57344, "num_bits": 8, "target_dtype": torch.float8_e5m2, "torch_dtype": torch.float8_e5m2, "storage_dtype": torch.float8_e5m2, "is_unsigned": False, "is_integer": False},
-    "float8_e4m3fnuz": {"min": -240, "max": 240, "num_bits": 8, "target_dtype": CustomDtype.FP8, "torch_dtype": torch.float8_e4m3fnuz, "storage_dtype": torch.float8_e4m3fnuz, "is_unsigned": False, "is_integer": False},
-    "float8_e5m2fnuz": {"min": -57344, "max": 57344, "num_bits": 8, "target_dtype": CustomDtype.FP8, "torch_dtype": torch.float8_e5m2fnuz, "storage_dtype": torch.float8_e5m2fnuz, "is_unsigned": False, "is_integer": False},
-}
-
-quantized_matmul_dtypes = ("int8", "int6", "int4", "int2", "float8_e4m3fn", "float8_e5m2")
-if devices.backend in {"cpu", "openvino"}:
-    quantized_matmul_dtypes += ("float8_e4m3fnuz", "float8_e5m2fnuz")
-
-linear_types = ("Linear",)
-conv_types = ("Conv1d", "Conv2d", "Conv3d")
-conv_transpose_types = ("ConvTranspose1d", "ConvTranspose2d", "ConvTranspose3d")
-allowed_types = linear_types + conv_types + conv_transpose_types
-
-
-class QuantizationMethod(str, Enum):
-    SDNQ = "sdnq"
-
-
-def sdnq_quantize_layer(layer, weights_dtype="int8", torch_dtype=None, group_size=0, quant_conv=False, use_quantized_matmul=False, param_name=None, pre_mode=False): # pylint: disable=unused-argument
-    layer_class_name = layer.__class__.__name__
-    if layer_class_name in allowed_types:
-        is_conv_type = False
-        is_conv_transpose_type = False
-        is_linear_type = False
-        use_tensorwise_fp8_matmul = False
-        result_shape = None
-        if torch_dtype is None:
-            torch_dtype = devices.dtype
-
-        if layer_class_name in conv_types:
-            if not quant_conv:
-                return layer
-            is_conv_type = True
-            reduction_axes = 1
-            output_channel_size, channel_size = layer.weight.shape[:2]
-            use_quantized_matmul = False
-            if dtype_dict[weights_dtype]["num_bits"] < 4:
-                weights_dtype = "uint4"
-        elif layer_class_name in conv_transpose_types:
-            if not quant_conv:
-                return layer
-            is_conv_transpose_type = True
-            reduction_axes = 0
-            channel_size, output_channel_size = layer.weight.shape[:2]
-            use_quantized_matmul = False
-            if dtype_dict[weights_dtype]["num_bits"] < 4:
-                weights_dtype = "uint4"
-        else:
-            is_linear_type = True
-            reduction_axes = -1
-            output_channel_size, channel_size = layer.weight.shape
-            if use_quantized_matmul:
-                use_quantized_matmul = weights_dtype in quantized_matmul_dtypes and channel_size >= 32 and output_channel_size >= 32
-                if use_quantized_matmul and not dtype_dict[weights_dtype]["is_integer"]:
-                    use_quantized_matmul = output_channel_size % 16 == 0 and channel_size % 16 == 0
-                    use_tensorwise_fp8_matmul = torch_version < 2.5 or devices.backend in {"cpu", "openvino"} or (devices.backend == "cuda" and sys.platform == "win32" and torch_version <= 2.7 and torch.cuda.get_device_capability(devices.device) == (8,9))
-
-        if group_size == 0:
-            if is_linear_type:
-                if dtype_dict[weights_dtype]["num_bits"] < 6:
-                    group_size = 2 ** (2 + dtype_dict[weights_dtype]["num_bits"])
-            else:
-                if dtype_dict[weights_dtype]["num_bits"] < 8:
-                    group_size = 2 ** (1 + dtype_dict[weights_dtype]["num_bits"])
-
-        if not use_quantized_matmul and group_size > 0:
-            if group_size >= channel_size:
-                group_size = channel_size
-                num_of_groups = 1
-            else:
-                num_of_groups = channel_size // group_size
-                while channel_size % group_size != 0: # find something divisible
-                    num_of_groups -= 1
-                    if num_of_groups <= 1:
-                        group_size = channel_size
-                        num_of_groups = 1
-                        break
-                    group_size = channel_size / num_of_groups
-            group_size = int(group_size)
-            num_of_groups = int(num_of_groups)
-
-            if num_of_groups > 1:
-                result_shape = layer.weight.shape
-                new_shape = list(result_shape)
-                if is_conv_type:
-                    # output_channel_size, channel_size, X, X
-                    # output_channel_size, num_of_groups, group_size, X, X
-                    new_shape[1] = group_size
-                    new_shape.insert(1, num_of_groups)
-                    reduction_axes = 2
-                elif is_conv_transpose_type:
-                    #channel_size, output_channel_size, X, X
-                    #num_of_groups, group_size, output_channel_size, X, X
-                    new_shape[0] = group_size
-                    new_shape.insert(0, num_of_groups)
-                    reduction_axes = 1
-                elif is_linear_type:
-                    # output_channel_size, channel_size
-                    # output_channel_size, num_of_groups, group_size
-                    last_dim_index = layer.weight.ndim
-                    new_shape[last_dim_index - 1 : last_dim_index] = (num_of_groups, group_size)
-                layer.weight.data = layer.weight.reshape(new_shape)
-
-        layer.weight.requires_grad = False
-        if shared.opts.diffusers_offload_mode in {"none", "model"}:
-            return_device = devices.device
-        elif pre_mode:
-            if shared.opts.device_map == "gpu":
-                return_device = devices.device
-            elif shared.opts.sdnq_quantize_with_gpu:
-                return_device = devices.cpu
-            else:
-                return_device = layer.weight.device
-        else:
-            return_device = layer.weight.device
-        if not pre_mode:
-            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)
-        else:
-            scale = get_scale_symmetric(layer.weight, reduction_axes, weights_dtype)
-            zero_point = None
-        layer.weight.data = quantize_weight(layer.weight, scale, zero_point, weights_dtype)
-
-        if not shared.opts.sdnq_decompress_fp32 and not (use_quantized_matmul and not dtype_dict[weights_dtype]["is_integer"] and not use_tensorwise_fp8_matmul):
-            scale = scale.to(torch_dtype)
-            if zero_point is not None:
-                zero_point = zero_point.to(torch_dtype)
-
-        if use_quantized_matmul:
-            scale = scale.transpose(0,1)
-            if dtype_dict[weights_dtype]["num_bits"] == 8:
-                layer.weight.data = layer.weight.transpose(0,1)
-            if not dtype_dict[weights_dtype]["is_integer"]:
-                stride = layer.weight.stride()
-                if stride[0] > stride[1] and stride[1] == 1:
-                    layer.weight.data = layer.weight.t().contiguous().t()
-                if not use_tensorwise_fp8_matmul:
-                    scale = scale.to(torch.float32)
-
-        layer.sdnq_decompressor = decompressor_dict[weights_dtype](
-            scale=scale,
-            zero_point=zero_point,
-            compressed_weight_shape=layer.weight.shape,
-            result_dtype=torch_dtype,
-            result_shape=result_shape,
-            weights_dtype=weights_dtype,
-            use_quantized_matmul=use_quantized_matmul,
-        )
-        layer.weight.data = layer.sdnq_decompressor.pack_weight(layer.weight).to(return_device)
-        layer.sdnq_decompressor = layer.sdnq_decompressor.to(return_device)
-
-        if is_linear_type:
-            if use_quantized_matmul:
-                if dtype_dict[weights_dtype]["is_integer"]:
-                    layer.forward = quantized_linear_forward_int8_matmul
-                else:
-                    if use_tensorwise_fp8_matmul:
-                        layer.forward = quantized_linear_forward_fp8_matmul_tensorwise
-                    else:
-                        layer.forward = quantized_linear_forward_fp8_matmul
-            else:
-                layer.forward = quantized_linear_forward
-        elif is_conv_type:
-            layer.forward = quantized_conv_forward
-        elif is_conv_transpose_type:
-            if layer_class_name.endswith("1d"):
-                layer.forward = quantized_conv_transpose_1d_forward
-            elif layer_class_name.endswith("2d"):
-                layer.forward = quantized_conv_transpose_2d_forward
-            elif layer_class_name.endswith("3d"):
-                layer.forward = quantized_conv_transpose_3d_forward
-        layer.forward = layer.forward.__get__(layer, layer.__class__)
-        devices.torch_gc(force=False, reason=f"SDNQ param_name: {param_name}")
-    return layer
-
-
-def apply_sdnq_to_module(model, weights_dtype="int8", torch_dtype=None, group_size=0, quant_conv=False, use_quantized_matmul=False, param_name=None):
-    has_children = list(model.children())
-    if not has_children:
-        return model
-    for module_param_name, module in model.named_children():
-        if hasattr(module, "weight") and module.weight is not None:
-            module = sdnq_quantize_layer(
-                module,
-                weights_dtype=weights_dtype,
-                torch_dtype=torch_dtype,
-                group_size=group_size,
-                quant_conv=quant_conv,
-                use_quantized_matmul=use_quantized_matmul,
-                param_name=module_param_name,
-            )
-        module = apply_sdnq_to_module(
-                module,
-                weights_dtype=weights_dtype,
-                torch_dtype=torch_dtype,
-                group_size=group_size,
-                quant_conv=quant_conv,
-                use_quantized_matmul=use_quantized_matmul,
-                param_name=module_param_name,
-            )
-    return model
-
-
-def get_scale_asymmetric(weight: torch.FloatTensor, reduction_axes: List[int], weights_dtype: str) -> Tuple[torch.FloatTensor, torch.FloatTensor]:
-    zero_point = torch.amin(weight, dim=reduction_axes, keepdims=True)
-    scale = torch.amax(weight, dim=reduction_axes, keepdims=True).sub_(zero_point).div_(dtype_dict[weights_dtype]["max"] - dtype_dict[weights_dtype]["min"])
-    eps = torch.finfo(scale.dtype).eps # prevent divison by 0
-    scale = torch.where(torch.abs(scale) < eps, eps, scale)
-    if dtype_dict[weights_dtype]["min"] != 0:
-        zero_point.sub_(torch.mul(scale, dtype_dict[weights_dtype]["min"]))
-    return scale, zero_point
-
-
-def get_scale_symmetric(weight: torch.FloatTensor, reduction_axes: List[int], weights_dtype: str) -> torch.FloatTensor:
-    abs_min_values = torch.amin(weight, dim=reduction_axes, keepdims=True).abs_()
-    max_values = torch.amax(weight, dim=reduction_axes, keepdims=True)
-    scale = torch.where(abs_min_values >= max_values, abs_min_values, -max_values).div_(dtype_dict[weights_dtype]["max"])
-    eps = torch.finfo(scale.dtype).eps # prevent divison by 0
-    scale = torch.where(torch.abs(scale) < eps, eps, scale)
-    return scale
-
-
-def quantize_weight(weight: torch.FloatTensor, scale: torch.FloatTensor, zero_point: torch.FloatTensor, weights_dtype: str) -> torch.ByteTensor:
-    if zero_point is not None:
-        compressed_weight = torch.sub(weight, zero_point).div_(scale)
-    else:
-        compressed_weight = torch.div(weight, scale)
-    if dtype_dict[weights_dtype]["is_integer"]:
-        compressed_weight.round_()
-    compressed_weight = compressed_weight.clamp_(dtype_dict[weights_dtype]["min"], dtype_dict[weights_dtype]["max"]).to(dtype_dict[weights_dtype]["torch_dtype"])
-    return compressed_weight
-
-
-def decompress_asymmetric(input: torch.Tensor, scale: torch.Tensor, zero_point: torch.Tensor, dtype: torch.dtype, result_shape: torch.Size) -> torch.Tensor:
-    result = torch.addcmul(zero_point, input.to(dtype=scale.dtype), scale).to(dtype=dtype)
-    if result_shape is not None:
-        result = result.reshape(result_shape)
-    return result
-
-
-def decompress_symmetric(input: torch.Tensor, scale: torch.Tensor, dtype: torch.dtype, result_shape: torch.Size, skip_quantized_matmul: bool = False) -> torch.Tensor:
-    if skip_quantized_matmul:
-        result = input.transpose(0,1).to(dtype=scale.dtype).mul_(scale.transpose(0,1)).to(dtype=dtype)
-    else:
-        result = input.to(dtype=scale.dtype).mul_(scale).to(dtype=dtype)
-    if result_shape is not None:
-        result = result.reshape(result_shape)
-    return result
-
-
-def decompress_packed_int_asymmetric(input: torch.Tensor, scale: torch.Tensor, zero_point: torch.Tensor, shape: torch.Size, dtype: torch.dtype, result_shape: torch.Size, weights_dtype: str) -> torch.Tensor:
-    return decompress_asymmetric(packed_int_function_dict[weights_dtype]["unpack"](input, shape), scale, zero_point, dtype, result_shape)
-
-
-def decompress_packed_int_symmetric(input: torch.Tensor, scale: torch.Tensor, shape: torch.Size, dtype: torch.dtype, result_shape: torch.Size, weights_dtype: str, skip_quantized_matmul: bool = False) -> torch.Tensor:
-    if skip_quantized_matmul:
-        return decompress_symmetric(unpack_int_symetric(input, shape, weights_dtype, dtype=scale.dtype), scale.transpose(0,1), dtype, result_shape)
-    else:
-        return decompress_symmetric(unpack_int_symetric(input, shape, weights_dtype, dtype=scale.dtype), scale, dtype, result_shape)
-
-
-def pack_int_symetric(tensor: torch.Tensor, weights_dtype: str) -> torch.Tensor:
-    return packed_int_function_dict[weights_dtype]["pack"](tensor.to(dtype=dtype_dict[weights_dtype]["torch_dtype"]).sub_(dtype_dict[weights_dtype]["min"]).to(dtype=dtype_dict[weights_dtype]["storage_dtype"]))
-
-
-def unpack_int_symetric(packed_tensor: torch.Tensor, shape: torch.Size, weights_dtype: str, dtype: Optional[torch.dtype] = None, transpose: Optional[bool] = False) -> torch.Tensor:
-    if dtype is None:
-        dtype = dtype_dict[weights_dtype]["torch_dtype"]
-    result = packed_int_function_dict[weights_dtype]["unpack"](packed_tensor, shape).to(dtype=dtype).add_(dtype_dict[weights_dtype]["min"])
-    if transpose:
-        result = result.transpose(0,1)
-    return result
-
-
-def pack_uint6(tensor: torch.Tensor) -> torch.Tensor:
-    if tensor.dtype != torch.uint8:
-        raise RuntimeError(f"Invalid tensor dtype {tensor.type}. torch.uint8 type is supported.")
-    packed_tensor = tensor.contiguous().reshape(-1, 4)
-    packed_tensor = torch.stack(
-        (
-            torch.bitwise_or(packed_tensor[:, 0], torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 3], 2), 192)),
-            torch.bitwise_or(packed_tensor[:, 1], torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 3], 4), 192)),
-            torch.bitwise_or(packed_tensor[:, 2], torch.bitwise_left_shift(packed_tensor[:, 3], 6)),
-        ),
-        dim=-1
-    )
-    return packed_tensor
-
-
-def pack_uint4(tensor: torch.Tensor) -> torch.Tensor:
-    if tensor.dtype != torch.uint8:
-        raise RuntimeError(f"Invalid tensor dtype {tensor.type}. torch.uint8 type is supported.")
-    packed_tensor = tensor.contiguous().reshape(-1, 2)
-    packed_tensor = torch.bitwise_or(packed_tensor[:, 0], torch.bitwise_left_shift(packed_tensor[:, 1], 4))
-    return packed_tensor
-
-
-def pack_uint2(tensor: torch.Tensor) -> torch.Tensor:
-    if tensor.dtype != torch.uint8:
-        raise RuntimeError(f"Invalid tensor dtype {tensor.type}. torch.uint8 type is supported.")
-    packed_tensor = tensor.contiguous().reshape(-1, 4)
-    packed_tensor = torch.bitwise_or(
-        torch.bitwise_or(packed_tensor[:, 0], torch.bitwise_left_shift(packed_tensor[:, 1], 2)),
-        torch.bitwise_or(torch.bitwise_left_shift(packed_tensor[:, 2], 4), torch.bitwise_left_shift(packed_tensor[:, 3], 6)),
-    )
-    return packed_tensor
-
-
-def unpack_uint6(packed_tensor: torch.Tensor, shape: torch.Size) -> torch.Tensor:
-    result = torch.stack(
-        (
-            torch.bitwise_and(packed_tensor[:, 0], 63),
-            torch.bitwise_and(packed_tensor[:, 1], 63),
-            torch.bitwise_and(packed_tensor[:, 2], 63),
-            torch.bitwise_or(
-                torch.bitwise_or(
-                    torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 0], 2), 48),
-                    torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 1], 4), 12),
-                ),
-                torch.bitwise_right_shift(packed_tensor[:, 2], 6)
-            )
-        ),
-        dim=-1
-    ).reshape(shape)
-    return result
-
-
-def unpack_uint4(packed_tensor: torch.Tensor, shape: torch.Size) -> torch.Tensor:
-    result = torch.stack((torch.bitwise_and(packed_tensor, 15), torch.bitwise_right_shift(packed_tensor, 4)), dim=-1).reshape(shape)
-    return result
-
-
-def unpack_uint2(packed_tensor: torch.Tensor, shape: torch.Size) -> torch.Tensor:
-    result = torch.stack(
-        (
-            torch.bitwise_and(packed_tensor, 3),
-            torch.bitwise_and(torch.bitwise_right_shift(packed_tensor, 2), 3),
-            torch.bitwise_and(torch.bitwise_right_shift(packed_tensor, 4), 3),
-            torch.bitwise_right_shift(packed_tensor, 6),
-        ),
-        dim=-1
-    ).reshape(shape)
-    return result
-
-
-def quantize_fp8_matmul_input(input: torch.FloatTensor) -> Tuple[torch.FloatTensor, torch.FloatTensor]:
-    input = input.flatten(0,-2).contiguous()
-    input_scale = torch.div(input.abs().amax(dim=-1, keepdims=True), 448)
-    input = torch.div(input, input_scale).clamp_(-448, 448).to(torch.float8_e4m3fn)
-    input_scale = input_scale.to(torch.float32)
-    return input, input_scale
-
-
-def quantize_fp8_matmul_input_tensorwise(input: torch.FloatTensor, scale: torch.FloatTensor) -> Tuple[torch.ByteTensor, torch.FloatTensor]:
-    input = input.flatten(0,-2).contiguous()
-    input_scale = torch.div(input.abs().amax(dim=-1, keepdims=True), 448)
-    input = torch.div(input, input_scale).clamp_(-448, 448).to(torch.float8_e4m3fn)
-    scale = torch.mul(input_scale, scale)
-    if scale.dtype == torch.float16: # fp16 will overflow
-        scale = scale.to(dtype=torch.float32)
-    return input, scale
-
-
-def quantize_int8_matmul_input(input: torch.FloatTensor, scale: torch.FloatTensor) -> Tuple[torch.ByteTensor, torch.FloatTensor]:
-    input = input.flatten(0,-2).contiguous()
-    input_scale = torch.div(input.abs().amax(dim=-1, keepdims=True), 127)
-    input = torch.div(input, input_scale).round_().clamp_(-128, 127).to(torch.int8)
-    scale = torch.mul(input_scale, scale)
-    if scale.dtype == torch.float16: # fp16 will overflow
-        scale = scale.to(dtype=torch.float32)
-    return input, scale
-
-
-def fp8_matmul(
-    input: torch.FloatTensor,
-    weight: torch.Tensor,
-    bias: torch.FloatTensor,
-    scale: torch.FloatTensor,
-) -> torch.FloatTensor:
-    return_dtype = input.dtype
-    output_shape = list(input.shape)
-    output_shape[-1] = weight.shape[-1]
-    input, input_scale = quantize_fp8_matmul_input(input)
-    return torch._scaled_mm(input, weight, scale_a=input_scale, scale_b=scale, bias=bias, out_dtype=return_dtype).reshape(output_shape)
-
-
-# sm89 doesn't support row wise scale in Windows
-def fp8_matmul_tensorwise(
-    input: torch.FloatTensor,
-    weight: torch.Tensor,
-    bias: torch.FloatTensor,
-    scale: torch.FloatTensor,
-) -> torch.FloatTensor:
-    return_dtype = input.dtype
-    output_shape = list(input.shape)
-    output_shape[-1] = weight.shape[-1]
-    dummy_input_scale = torch.ones(1, device=input.device, dtype=torch.float32)
-    input, scale = quantize_fp8_matmul_input_tensorwise(input, scale)
-    result = decompress_symmetric(torch._scaled_mm(input, weight, scale_a=dummy_input_scale, scale_b=dummy_input_scale, bias=None, out_dtype=scale.dtype), scale, return_dtype, output_shape)
-    if bias is not None:
-        result.add_(bias)
-    return result
-
-
-def int8_matmul(
-    input: torch.FloatTensor,
-    weight: torch.Tensor,
-    bias: torch.FloatTensor,
-    scale: torch.FloatTensor,
-    compressed_weight_shape: torch.Size,
-    weights_dtype: str,
-) -> torch.FloatTensor:
-    if compressed_weight_shape is not None:
-        weight = unpack_int_symetric(weight, compressed_weight_shape, weights_dtype, dtype=torch.int8, transpose=True)
-    return_dtype = input.dtype
-    output_shape = list(input.shape)
-    output_shape[-1] = weight.shape[-1]
-    input, scale = quantize_int8_matmul_input(input, scale)
-    result = decompress_symmetric(torch._int_mm(input, weight), scale, return_dtype, output_shape)
-    if bias is not None:
-        result.add_(bias)
-    return result
-
-
-def quantized_linear_forward_fp8_matmul(self, input: torch.FloatTensor) -> torch.FloatTensor:
-    if torch.numel(input) / input.shape[-1] < 32:
-        return torch.nn.functional.linear(input, self.sdnq_decompressor(self.weight, skip_quantized_matmul=True), self.bias)
-    return fp8_matmul(input, self.weight, self.bias, self.sdnq_decompressor.scale)
-
-
-def quantized_linear_forward_fp8_matmul_tensorwise(self, input: torch.FloatTensor) -> torch.FloatTensor:
-    if torch.numel(input) / input.shape[-1] < 32:
-        return torch.nn.functional.linear(input, self.sdnq_decompressor(self.weight, skip_quantized_matmul=True), self.bias)
-    return fp8_matmul_tensorwise(input, self.weight, self.bias, self.sdnq_decompressor.scale)
-
-
-def quantized_linear_forward_int8_matmul(self, input: torch.FloatTensor) -> torch.FloatTensor:
-    if torch.numel(input) / input.shape[-1] < 32:
-        return torch.nn.functional.linear(input, self.sdnq_decompressor(self.weight, skip_quantized_matmul=True), self.bias)
-    return int8_matmul(input, self.weight, self.bias, self.sdnq_decompressor.scale, getattr(self.sdnq_decompressor, "compressed_weight_shape", None), self.sdnq_decompressor.weights_dtype)
-
-
-def quantized_linear_forward(self, input: torch.FloatTensor) -> torch.FloatTensor:
-    return torch.nn.functional.linear(input, self.sdnq_decompressor(self.weight), self.bias)
-
-
-def quantized_conv_forward(self, input) -> torch.FloatTensor:
-    return self._conv_forward(input, self.sdnq_decompressor(self.weight), self.bias)
-
-
-def quantized_conv_transpose_1d_forward(self, input: torch.FloatTensor, output_size: Optional[list[int]] = None) -> torch.FloatTensor:
-    output_padding = self._output_padding(input, output_size, self.stride, self.padding, self.kernel_size, 1, self.dilation)
-    return torch.nn.functional.conv_transpose1d(input, self.sdnq_decompressor(self.weight), self.bias, self.stride, self.padding, output_padding, self.groups, self.dilation)
-
-
-def quantized_conv_transpose_2d_forward(self, input: torch.FloatTensor, output_size: Optional[list[int]] = None) -> torch.FloatTensor:
-    output_padding = self._output_padding(input, output_size, self.stride, self.padding, self.kernel_size, 2, self.dilation)
-    return torch.nn.functional.conv_transpose2d(input, self.sdnq_decompressor(self.weight), self.bias, self.stride, self.padding, output_padding, self.groups, self.dilation)
-
-
-def quantized_conv_transpose_3d_forward(self, input: torch.FloatTensor, output_size: Optional[list[int]] = None) -> torch.FloatTensor:
-    output_padding = self._output_padding(input, output_size, self.stride, self.padding, self.kernel_size, 3, self.dilation)
-    return torch.nn.functional.conv_transpose3d(input, self.sdnq_decompressor(self.weight), self.bias, self.stride, self.padding, output_padding, self.groups, self.dilation)
-
-
-class AsymmetricWeightsDecompressor(torch.nn.Module):
-    def __init__(
-        self,
-        scale: torch.Tensor,
-        zero_point: torch.Tensor,
-        result_dtype: torch.dtype,
-        result_shape: torch.Size,
-        weights_dtype: str,
-        **kwargs,
-    ):
-        super().__init__()
-        self.weights_dtype = weights_dtype
-        self.use_quantized_matmul = False
-        self.result_dtype = result_dtype
-        self.result_shape = result_shape
-        self.register_buffer("scale", scale)
-        self.register_buffer("zero_point", zero_point)
-
-    def pack_weight(self, weight: torch.Tensor) -> torch.Tensor:
-        return weight.to(dtype=dtype_dict[self.weights_dtype]["torch_dtype"])
-
-    def forward(self, weight, **kwargs):
-        return decompress_asymmetric_compiled(weight, self.scale, self.zero_point, self.result_dtype, self.result_shape)
-
-
-class SymmetricWeightsDecompressor(torch.nn.Module):
-    def __init__(
-        self,
-        scale: torch.Tensor,
-        result_dtype: torch.dtype,
-        result_shape: torch.Size,
-        weights_dtype: str,
-        use_quantized_matmul: bool = False,
-        **kwargs,
-    ):
-        super().__init__()
-        self.weights_dtype = weights_dtype
-        self.use_quantized_matmul = use_quantized_matmul
-        self.result_dtype = result_dtype
-        self.result_shape = result_shape
-        self.register_buffer("scale", scale)
-
-    def pack_weight(self, weight: torch.Tensor) -> torch.Tensor:
-        return weight.to(dtype=dtype_dict[self.weights_dtype]["torch_dtype"])
-
-    def forward(self, weight, skip_quantized_matmul=False, **kwargs):
-        return decompress_symmetric_compiled(weight, self.scale, self.result_dtype, self.result_shape, skip_quantized_matmul=skip_quantized_matmul)
-
-
-class PackedINTAsymmetricWeightsDecompressor(torch.nn.Module):
-    def __init__(
-        self,
-        scale: torch.Tensor,
-        zero_point: torch.Tensor,
-        compressed_weight_shape: torch.Size,
-        result_dtype: torch.dtype,
-        result_shape: torch.Size,
-        weights_dtype: str,
-        **kwargs,
-    ):
-        super().__init__()
-        self.weights_dtype = weights_dtype
-        self.use_quantized_matmul = False
-        self.compressed_weight_shape = compressed_weight_shape
-        self.result_dtype = result_dtype
-        self.result_shape = result_shape
-        self.register_buffer("scale", scale)
-        self.register_buffer("zero_point", zero_point)
-
-    def pack_weight(self, weight: torch.Tensor) -> torch.Tensor:
-        return packed_int_function_dict[self.weights_dtype]["pack"](weight.to(dtype=dtype_dict[self.weights_dtype]["torch_dtype"]))
-
-    def forward(self, weight, **kwargs):
-        return decompress_packed_int_asymmetric_compiled(weight, self.scale, self.zero_point, self.compressed_weight_shape, self.result_dtype, self.result_shape, self.weights_dtype)
-
-
-class PackedINTSymmetricWeightsDecompressor(torch.nn.Module):
-    def __init__(
-        self,
-        scale: torch.Tensor,
-        compressed_weight_shape: torch.Size,
-        result_dtype: torch.dtype,
-        result_shape: torch.Size,
-        weights_dtype: str,
-        use_quantized_matmul: bool = False,
-        **kwargs,
-    ):
-        super().__init__()
-        self.weights_dtype = weights_dtype
-        self.use_quantized_matmul = use_quantized_matmul
-        self.compressed_weight_shape = compressed_weight_shape
-        self.result_dtype = result_dtype
-        self.result_shape = result_shape
-        self.register_buffer("scale", scale)
-
-    def pack_weight(self, weight: torch.Tensor) -> torch.Tensor:
-        return pack_int_symetric(weight, self.weights_dtype)
-
-    def forward(self, weight, skip_quantized_matmul=False, **kwargs):
-        return decompress_packed_int_symmetric_compiled(weight, self.scale, self.compressed_weight_shape, self.result_dtype, self.result_shape, self.weights_dtype, skip_quantized_matmul=skip_quantized_matmul)
-
-
-decompressor_dict = {
-    "int8": SymmetricWeightsDecompressor,
-    "uint8": AsymmetricWeightsDecompressor,
-    "int6": PackedINTSymmetricWeightsDecompressor,
-    "uint6": PackedINTAsymmetricWeightsDecompressor,
-    "int4": PackedINTSymmetricWeightsDecompressor,
-    "uint4": PackedINTAsymmetricWeightsDecompressor,
-    "int2": PackedINTSymmetricWeightsDecompressor,
-    "uint2": PackedINTAsymmetricWeightsDecompressor,
-    "uint1": AsymmetricWeightsDecompressor,
-    "float8_e4m3fn": SymmetricWeightsDecompressor,
-    "float8_e4m3fnuz": SymmetricWeightsDecompressor,
-    "float8_e5m2": SymmetricWeightsDecompressor,
-    "float8_e5m2fnuz": SymmetricWeightsDecompressor,
-}
-
-
-packed_int_function_dict = {
-    "int6": {"pack": pack_uint6, "unpack": unpack_uint6},
-    "uint6": {"pack": pack_uint6, "unpack": unpack_uint6},
-    "int4": {"pack": pack_uint4, "unpack": unpack_uint4},
-    "uint4": {"pack": pack_uint4, "unpack": unpack_uint4},
-    "int2": {"pack": pack_uint2, "unpack": unpack_uint2},
-    "uint2": {"pack": pack_uint2, "unpack": unpack_uint2},
-}
-
-
-class SDNQQuantizer(DiffusersQuantizer):
-    r"""
-    Diffusers Quantizer for SDNQ
-    """
-
-    requires_parameters_quantization = True
-    use_keep_in_fp32_modules = True
-    requires_calibration = False
-    required_packages = None
-    torch_dtype = None
-
-    def __init__(self, quantization_config, **kwargs): # pylint: disable=useless-parent-delegation
-        super().__init__(quantization_config, **kwargs)
-
-    def check_if_quantized_param(
-        self,
-        model,
-        param_value: "torch.Tensor",
-        param_name: str,
-        state_dict: Dict[str, Any],
-        **kwargs,
-    ):
-        if param_name.endswith(".weight"):
-            split_param_name = param_name.split(".")
-            if param_name not in self.modules_to_not_convert and not any(param in split_param_name for param in self.modules_to_not_convert):
-                layer_class_name = get_module_from_name(model, param_name)[0].__class__.__name__
-                if layer_class_name in allowed_types:
-                    if layer_class_name in conv_types or layer_class_name in conv_transpose_types:
-                        if self.quantization_config.quant_conv:
-                            return True
-                    else:
-                        return True
-        param_value.data = param_value.clone() # safetensors is unable to release the cpu memory without this
-        return False
-
-    def check_quantized_param(self, *args, **kwargs) -> bool:
-        """
-        needed for transformers compatibilty, returns self.check_if_quantized_param
-        """
-        return self.check_if_quantized_param(*args, **kwargs)
-
-    def create_quantized_param( # pylint: disable=arguments-differ
-        self,
-        model,
-        param_value: torch.FloatTensor,
-        param_name: str,
-        target_device: torch.device,
-        state_dict: Dict[str, Any], # pylint: disable=unused-argument
-        unexpected_keys: List[str], # pylint: disable=unused-argument
-        **kwargs,
-    ):
-        # load the model params to target_device first
-        layer, _ = get_module_from_name(model, param_name)
-        if shared.opts.sdnq_quantize_with_gpu:
-            if param_value.dtype == torch.float32 and devices.same_device(param_value.device, devices.device):
-                param_value = param_value.clone()
-            else:
-                param_value = param_value.to(devices.device).to(dtype=torch.float32)
-        else:
-            if param_value.dtype == torch.float32 and devices.same_device(param_value.device, target_device):
-                param_value = param_value.clone()
-            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,
-            torch_dtype=self.torch_dtype,
-            group_size=self.quantization_config.group_size,
-            quant_conv=self.quantization_config.quant_conv,
-            use_quantized_matmul=self.quantization_config.use_quantized_matmul,
-            param_name=param_name,
-            pre_mode=True,
-        )
-
-    def adjust_max_memory(self, max_memory: Dict[str, Union[int, str]]) -> Dict[str, Union[int, str]]:
-        max_memory = {key: val * 0.80 for key, val in max_memory.items()}
-        return max_memory
-
-    def adjust_target_dtype(self, target_dtype: torch.dtype) -> torch.dtype: # pylint: disable=unused-argument,arguments-renamed
-        return dtype_dict[self.quantization_config.weights_dtype]["target_dtype"]
-
-    def update_torch_dtype(self, torch_dtype: torch.dtype = None) -> torch.dtype:
-        if torch_dtype is None:
-            torch_dtype = devices.dtype
-        self.torch_dtype = torch_dtype
-        return torch_dtype
-
-    def _process_model_before_weight_loading( # pylint: disable=arguments-differ
-        self,
-        model,
-        device_map, # pylint: disable=unused-argument
-        keep_in_fp32_modules: List[str] = [],
-        **kwargs,
-    ):
-        model.config.quantization_config = self.quantization_config
-        self.modules_to_not_convert = self.quantization_config.modules_to_not_convert
-        if not isinstance(self.modules_to_not_convert, list):
-            self.modules_to_not_convert = [self.modules_to_not_convert]
-        if keep_in_fp32_modules is not None:
-            self.modules_to_not_convert.extend(keep_in_fp32_modules)
-
-    def _process_model_after_weight_loading(self, model, **kwargs):
-        if shared.opts.diffusers_offload_mode != "none":
-            model = model.to(devices.cpu)
-        devices.torch_gc(force=True)
-        return model
-
-    def get_cuda_warm_up_factor(self):
-        return 32 // dtype_dict[self.quantization_config.weights_dtype]["num_bits"]
-
-    def update_tp_plan(self, config):
-        """
-        needed for transformers compatibilty, no-op function
-        """
-        return config
-
-    def update_unexpected_keys(self, model, unexpected_keys: List[str], prefix: str) -> List[str]: # pylint: disable=unused-argument
-        """
-        needed for transformers compatibilty, no-op function
-        """
-        return unexpected_keys
-
-    def update_missing_keys_after_loading(self, model, missing_keys: List[str], prefix: str) -> List[str]: # pylint: disable=unused-argument
-        """
-        needed for transformers compatibilty, no-op function
-        """
-        return missing_keys
-
-    def update_expected_keys(self, model, expected_keys: List[str], loaded_keys: List[str]) -> List[str]: # pylint: disable=unused-argument
-        """
-        needed for transformers compatibilty, no-op function
-        """
-        return expected_keys
-
-    @property
-    def is_trainable(self):
-        return False
-
-    @property
-    def is_serializable(self):
-        return False
-
-
-@dataclass
-class SDNQConfig(QuantizationConfigMixin):
-    """
-    This is a wrapper class about all possible attributes and features that you can play with a model that has been
-    loaded using `sdnq`.
-
-    Args:
-        weights_dtype (`str`, *optional*, defaults to `"int8"`):
-            The target dtype for the weights after quantization. Supported values are:
-            ("int8", "uint8", "int6", "uint6", "int4", "uint4", "uint2", "uint1", "float8_e4m3fn", "float8_e4m3fnuz", "float8_e5m2", "float8_e5m2fnuz")
-       modules_to_not_convert (`list`, *optional*, default to `None`):
-            The list of modules to not quantize, useful for quantizing models that explicitly require to have some
-            modules left in their original precision (e.g. Whisper encoder, Llava encoder, Mixtral gate layers).
-    """
-
-    def __init__( # pylint: disable=super-init-not-called
-        self,
-        weights_dtype: str = "int8",
-        group_size: int = 0,
-        quant_conv: bool = False,
-        use_quantized_matmul: bool = False,
-        modules_to_not_convert: Optional[List[str]] = None,
-        **kwargs, # pylint: disable=unused-argument
-    ):
-        self.weights_dtype = weights_dtype
-        self.quant_method = QuantizationMethod.SDNQ
-        self.group_size = group_size
-        self.quant_conv = quant_conv
-        self.use_quantized_matmul = use_quantized_matmul
-        self.modules_to_not_convert = modules_to_not_convert
-        self.post_init()
-        self.is_integer = dtype_dict[self.weights_dtype]["is_integer"]
-
-    def post_init(self):
-        r"""
-        Safety checker that arguments are correct
-        """
-        accepted_weights = ["int8", "uint8", "int6", "uint6", "int4", "uint4", "uint2", "uint1", "float8_e4m3fn", "float8_e4m3fnuz", "float8_e5m2", "float8_e5m2fnuz"]
-        if self.weights_dtype not in accepted_weights:
-            raise ValueError(f"Only support weights in {accepted_weights} but found {self.weights_dtype}")
-
-
-class SDNQ_T5DenseGatedActDense(torch.nn.Module): # forward can't find what self is without creating a class
-    def __init__(self, T5DenseGatedActDense, dtype):
-        super().__init__()
-        self.wi_0 = T5DenseGatedActDense.wi_0
-        self.wi_1 = T5DenseGatedActDense.wi_1
-        self.wo = T5DenseGatedActDense.wo
-        self.dropout = T5DenseGatedActDense.dropout
-        self.act = T5DenseGatedActDense.act
-        self.torch_dtype = dtype
-
-    def forward(self, hidden_states):
-        hidden_gelu = self.act(self.wi_0(hidden_states))
-        hidden_linear = self.wi_1(hidden_states)
-        hidden_states = hidden_gelu * hidden_linear
-        hidden_states = self.dropout(hidden_states)
-        hidden_states = hidden_states.to(self.torch_dtype) # this line needs to be forced
-        hidden_states = self.wo(hidden_states)
-        return hidden_states
-
-
-if shared.opts.sdnq_decompress_compile:
-    try:
-        torch._dynamo.config.cache_size_limit = max(8192, torch._dynamo.config.cache_size_limit)
-        decompress_asymmetric_compiled = torch.compile(decompress_asymmetric, fullgraph=True)
-        decompress_symmetric_compiled = torch.compile(decompress_symmetric, fullgraph=True)
-        decompress_packed_int_asymmetric_compiled = torch.compile(decompress_packed_int_asymmetric, fullgraph=True)
-        decompress_packed_int_symmetric_compiled = torch.compile(decompress_packed_int_symmetric, fullgraph=True)
-        fp8_matmul = torch.compile(fp8_matmul, fullgraph=True)
-        fp8_matmul_tensorwise = torch.compile(fp8_matmul_tensorwise, fullgraph=True)
-        int8_matmul = torch.compile(int8_matmul, fullgraph=True)
-    except Exception as e:
-        shared.log.warning(f"Quantization: type=sdnq Decompress using torch.compile is not available: {e}")
-        decompress_asymmetric_compiled = decompress_asymmetric
-        decompress_symmetric_compiled = decompress_symmetric
-        decompress_packed_int_asymmetric_compiled = decompress_packed_int_asymmetric
-        decompress_packed_int_symmetric_compiled = decompress_packed_int_symmetric
-else:
-    decompress_asymmetric_compiled = decompress_asymmetric
-    decompress_symmetric_compiled = decompress_symmetric
-    decompress_packed_int_asymmetric_compiled = decompress_packed_int_asymmetric
-    decompress_packed_int_symmetric_compiled = decompress_packed_int_symmetric

modules/modeldata.py

@@ -56,6 +56,10 @@ def get_model_type(pipe):
         model_type = 'mochivideo'
     elif "Allegro" in name:
         model_type = 'allegrovideo'
+    elif "PixArtSigma" in name:
+        model_type = 'pixartsigma'
+    elif "PixArtAlpha" in name:
+        model_type = 'pixartalpha'
     else:
         model_type = name
     return model_type

modules/omnigen/transformer.py

@@ -99,6 +99,9 @@ class Phi3Transformer(Phi3Model):
 
         hidden_states = inputs_embeds
 
+        # create position embeddings to be shared across the decoder layers
+        position_embeddings = self.rotary_emb(hidden_states, position_ids)
+
         # decoder layers
         all_hidden_states = () if output_hidden_states else None
         all_self_attns = () if output_attentions else None
@@ -118,6 +121,7 @@ class Phi3Transformer(Phi3Model):
                     output_attentions,
                     use_cache,
                     cache_position,
+                    position_embeddings,
                 )
             else:
                 layer_outputs = decoder_layer(
@@ -128,6 +132,7 @@ class Phi3Transformer(Phi3Model):
                     output_attentions=output_attentions,
                     use_cache=use_cache,
                     cache_position=cache_position,
+                    position_embeddings=position_embeddings,
                 )
 
             hidden_states = layer_outputs[0]

modules/postprocess/gfpgan_model.py

@@ -72,8 +72,8 @@ def setup_model(dirname):
     except Exception:
         pass
     try:
-        install('basicsr', quiet=True)
-        install('gfpgan', quiet=True)
+        install('git+https://github.com/Disty0/BasicSR@2b6a12c28e0c81bfb13b7e984144f0b0f5461484', 'basicsr')
+        install('git+https://github.com/Disty0/GFPGAN@09b1190eabbc77e5f15c61fa7c38a2064b403e20', 'gfpgan')
         import gfpgan
         import facexlib
         import modules.detailer

modules/processing_vae.py

@@ -117,6 +117,7 @@ def full_vae_decode(latents, model):
     elif shared.opts.diffusers_offload_mode != "sequential":
         sd_models.move_model(model.vae, devices.device)
 
+    sd_models.set_vae_options(model, vae=None, op='decode')
     upcast = (model.vae.dtype == torch.float16) and (getattr(model.vae.config, 'force_upcast', False) or shared.opts.no_half_vae)
     if upcast:
         if hasattr(model, 'upcast_vae'): # this is done by diffusers automatically if output_type != 'latent'
@@ -125,10 +126,6 @@ def full_vae_decode(latents, model):
             model.vae.orig_dtype = model.vae.dtype
             model.vae = model.vae.to(dtype=torch.float32)
     latents = latents.to(devices.device)
-    if getattr(model.vae, "post_quant_conv", None) is not None:
-        latents = latents.to(next(iter(model.vae.post_quant_conv.parameters())).dtype)
-    else:
-        latents = latents.to(model.vae.dtype)
 
     # normalize latents
     latents_mean = model.vae.config.get("latents_mean", None)
@@ -144,6 +141,16 @@ def full_vae_decode(latents, model):
     if shift_factor:
         latents = latents + shift_factor
 
+    if getattr(model.vae, "post_quant_conv", None) is not None:
+        if getattr(model.vae.post_quant_conv, "bias", None) is not None:
+            latents = latents.to(model.vae.post_quant_conv.bias.dtype)
+        elif "VAE" in shared.opts.sdnq_quantize_weights:
+            latents = latents.to(devices.dtype_vae)
+        else:
+            latents = latents.to(next(iter(model.vae.post_quant_conv.parameters())).dtype)
+    else:
+        latents = latents.to(model.vae.dtype)
+
     log_debug(f'VAE config: {model.vae.config}')
     try:
         with devices.inference_context():
@@ -187,6 +194,7 @@ def full_vae_encode(image, model):
     vae_name = sd_vae.loaded_vae_file if sd_vae.loaded_vae_file is not None else "default"
     log_debug(f'Encode vae="{vae_name}" dtype={model.vae.dtype} upcast={model.vae.config.get("force_upcast", None)}')
 
+    sd_models.set_vae_options(model, vae=None, op='encode')
     upcast = (model.vae.dtype == torch.float16) and (getattr(model.vae.config, 'force_upcast', False) or shared.opts.no_half_vae)
     if upcast:
         if hasattr(model, 'upcast_vae'): # this is done by diffusers automatically if output_type != 'latent'
@@ -248,8 +256,8 @@ def vae_postprocess(tensor, model, output_type='np'):
                 if output_type == "pil":
                     images = model.numpy_to_pil(images)
             else:
-                import diffusers
-                model.image_processor = diffusers.image_processor.VaeImageProcessor()
+                from diffusers.image_processor import VaeImageProcessor
+                model.image_processor = VaeImageProcessor()
                 images = model.image_processor.postprocess(tensor, output_type=output_type)
         else:
             images = tensor if isinstance(tensor, list) or isinstance(tensor, np.ndarray) else [tensor]

modules/rocm.py

@@ -81,11 +81,9 @@ class Agent:
         self.blaslt_supported = os.path.exists(os.path.join(blaslt_tensile_libpath, f"Kernels.so-000-{name}.hsaco" if sys.platform == "win32" else f"extop_{name}.co"))
 
     def get_gfx_version(self) -> Union[str, None]:
-        if self.gfx_version >= 0x1200:
-            return "12.0.0"
-        elif self.gfx_version >= 0x1100:
+        if self.gfx_version >= 0x1101 and self.gfx_version < 0x1200:
             return "11.0.0"
-        elif self.gfx_version >= 0x1000:
+        elif self.gfx_version != 0x1030 and self.gfx_version >= 0x1000 and self.gfx_version < 0x1100:
             # gfx1010 users had to override gfx version to 10.3.0 in Linux
             # it is unknown whether overriding is needed in ZLUDA
             return "10.3.0"
@@ -206,7 +204,7 @@ else:
         if agent.gfx_version >= 0x1100 and os.environ.get("FLASH_ATTENTION_USE_TRITON_ROCM", "false").lower() != "true":
             # use the navi_rotary_fix fork because the original doesn't support rotary_emb for transformers
             # original: "git+https://github.com/ROCm/flash-attention@howiejay/navi_support"
-            default = "https://github.com/Disty0/flash-attention@navi_rotary_fix"
+            default = "git+https://github.com/Disty0/flash-attention@navi_rotary_fix"
         return os.environ.get("FLASH_ATTENTION_PACKAGE", default)
 
     is_wsl: bool = os.environ.get('WSL_DISTRO_NAME', 'unknown' if spawn('wslpath -w /') else None) is not None

modules/sd_models.py

@@ -92,7 +92,7 @@ def set_vae_options(sd_model, vae=None, op:str='model', quiet:bool=False):
         if shared.opts.diffusers_vae_upcast != 'default':
             sd_model.vae.config.force_upcast = True if shared.opts.diffusers_vae_upcast == 'true' else False
             shared.log.quiet(quiet, f'Setting {op}: component=VAE upcast={sd_model.vae.config.force_upcast}')
-        if shared.opts.no_half_vae:
+        if shared.opts.no_half_vae and op not in {'decode', 'encode'}:
             devices.dtype_vae = torch.float32
             sd_model.vae.to(devices.dtype_vae)
             shared.log.quiet(quiet, f'Setting {op}: component=VAE no-half=True')
@@ -105,11 +105,20 @@ def set_vae_options(sd_model, vae=None, op:str='model', quiet:bool=False):
     if hasattr(sd_model, "enable_vae_tiling"):
         if shared.opts.diffusers_vae_tiling:
             if hasattr(sd_model, 'vae') and hasattr(sd_model.vae, 'config') and hasattr(sd_model.vae.config, 'sample_size') and isinstance(sd_model.vae.config.sample_size, int):
+                if getattr(sd_model.vae, "tile_sample_min_size_backup", None) is None:
+                    sd_model.vae.tile_sample_min_size_backup = sd_model.vae.tile_sample_min_size
+                    sd_model.vae.tile_latent_min_size_backup = sd_model.vae.tile_latent_min_size
+                    sd_model.vae.tile_overlap_factor_backup = sd_model.vae.tile_overlap_factor
                 if shared.opts.diffusers_vae_tile_size > 0:
                     sd_model.vae.tile_sample_min_size = int(shared.opts.diffusers_vae_tile_size)
-                    sd_model.vae.tile_latent_min_size = int(sd_model.vae.config.sample_size / (2 ** (len(sd_model.vae.config.block_out_channels) - 1)))
+                    sd_model.vae.tile_latent_min_size = int(shared.opts.diffusers_vae_tile_size / (2 ** (len(sd_model.vae.config.block_out_channels) - 1)))
+                else:
+                    sd_model.vae.tile_sample_min_size = getattr(sd_model.vae, "tile_sample_min_size_backup", sd_model.vae.tile_sample_min_size)
+                    sd_model.vae.tile_latent_min_size = getattr(sd_model.vae, "tile_latent_min_size_backup", sd_model.vae.tile_latent_min_size)
                 if shared.opts.diffusers_vae_tile_overlap != 0.25:
                     sd_model.vae.tile_overlap_factor = float(shared.opts.diffusers_vae_tile_overlap)
+                else:
+                    sd_model.vae.tile_overlap_factor = getattr(sd_model.vae, "tile_overlap_factor_backup", sd_model.vae.tile_overlap_factor)
                 shared.log.quiet(quiet, f'Setting {op}: component=VAE tiling=True tile={sd_model.vae.tile_sample_min_size} overlap={sd_model.vae.tile_overlap_factor}')
             else:
                 shared.log.quiet(quiet, f'Setting {op}: component=VAE tiling=True')
@@ -899,48 +908,45 @@ def set_diffuser_pipe(pipe, new_pipe_type):
 def set_diffusers_attention(pipe, quiet:bool=False):
     import diffusers.models.attention_processor as p
 
-    def set_attn(pipe, attention):
+    def set_attn(pipe, attention, name:str=None, quiet:bool=False):
         if attention is None:
             return
-        if not hasattr(pipe, "_internal_dict"):
-            return
-        modules = [getattr(pipe, n, None) for n in pipe._internal_dict.keys()] # pylint: disable=protected-access
-        modules = [m for m in modules if isinstance(m, torch.nn.Module) and hasattr(m, "set_attn_processor")]
-        for module in modules:
-            if module.__class__.__name__ in ['SD3Transformer2DModel']:
-                module.set_attn_processor(p.JointAttnProcessor2_0())
-            elif module.__class__.__name__ in ['FluxTransformer2DModel']:
-                module.set_attn_processor(p.FluxAttnProcessor2_0())
-            elif module.__class__.__name__ in ['HunyuanDiT2DModel']:
-                module.set_attn_processor(p.HunyuanAttnProcessor2_0())
-            elif module.__class__.__name__ in ['AuraFlowTransformer2DModel']:
-                module.set_attn_processor(p.AuraFlowAttnProcessor2_0())
-            elif 'KandinskyCombinedPipeline' in pipe.__class__.__name__:
-                pass
-            elif 'Transformer' in module.__class__.__name__:
-                pass # unknown transformer so probably dont want to force attention processor
-            else:
-                module.set_attn_processor(attention)
+        # other models uses their own attention processor
+        if pipe.__class__.__name__.startswith("StableDiffusion") and hasattr(pipe, "unet"):
+            pipe.unet.set_attn_processor(attention)
+        elif not quiet:
+            shared.log.warning(f"Attention: {name if name is not None else attention.__class__.__name__} is not compatible with {pipe.__class__.__name__}")
 
     # if hasattr(pipe, 'pipe'):
     #    set_diffusers_attention(pipe.pipe)
 
-    if 'ControlNet' in pipe.__class__.__name__: # do not replace attention in ControlNet pipelines
+    if 'ControlNet' in pipe.__class__.__name__ or not (pipe.__class__.__name__.startswith("StableDiffusion") and hasattr(pipe, "unet")):
+        if shared.opts.cross_attention_optimization not in {"Scaled-Dot-Product", "Disabled"}:
+            shared.log.warning(f"Attention: {shared.opts.cross_attention_optimization} is not compatible with {pipe.__class__.__name__}")
+        else:
+            pipe.current_attn_name = shared.opts.cross_attention_optimization
         return
+
     shared.log.quiet(quiet, f'Setting model: attention="{shared.opts.cross_attention_optimization}"')
     if shared.opts.cross_attention_optimization == "Disabled":
         pass # do nothing
     elif shared.opts.cross_attention_optimization == "Scaled-Dot-Product": # The default set by Diffusers
-        set_attn(pipe, p.AttnProcessor2_0())
-    elif shared.opts.cross_attention_optimization == "xFormers" and hasattr(pipe, 'enable_xformers_memory_efficient_attention'):
-        pipe.enable_xformers_memory_efficient_attention()
-    elif shared.opts.cross_attention_optimization == "Split attention" and hasattr(pipe, "enable_attention_slicing"):
-        pipe.enable_attention_slicing()
+        set_attn(pipe, p.AttnProcessor2_0(), name="Scaled-Dot-Product", quiet=True)
+    elif shared.opts.cross_attention_optimization == "xFormers":
+        if hasattr(pipe, 'enable_xformers_memory_efficient_attention'):
+            pipe.enable_xformers_memory_efficient_attention()
+        else:
+            shared.log.warning(f"Attention: xFormers is not compatible with {pipe.__class__.__name__}")
+    elif shared.opts.cross_attention_optimization == "Split attention":
+        if hasattr(pipe, "enable_attention_slicing"):
+            pipe.enable_attention_slicing()
+        else:
+            shared.log.warning(f"Attention: Split attention is not compatible with {pipe.__class__.__name__}")
     elif shared.opts.cross_attention_optimization == "Batch matrix-matrix":
-        set_attn(pipe, p.AttnProcessor())
+        set_attn(pipe, p.AttnProcessor(), name="Batch matrix-matrix")
     elif shared.opts.cross_attention_optimization == "Dynamic Attention BMM":
         from modules.sd_hijack_dynamic_atten import DynamicAttnProcessorBMM
-        set_attn(pipe, DynamicAttnProcessorBMM())
+        set_attn(pipe, DynamicAttnProcessorBMM(), name="Dynamic Attention BMM")
 
     pipe.current_attn_name = shared.opts.cross_attention_optimization
 

modules/sd_models_utils.py

@@ -155,6 +155,8 @@ def apply_function_to_model(sd_model, function, options, op=None):
         if hasattr(sd_model, 'transformer') and hasattr(sd_model.transformer, 'config'):
             sd_model.transformer = function(sd_model.transformer, op="transformer", sd_model=sd_model)
     if "Model" in options:
+        if hasattr(sd_model, 'model') and (hasattr(sd_model.model, 'config') or isinstance(sd_model.model, torch.nn.Module)):
+            sd_model.model = function(sd_model.model, op="model", sd_model=sd_model)
         if hasattr(sd_model, 'unet') and hasattr(sd_model.unet, 'config'):
             sd_model.unet = function(sd_model.unet, op="unet", sd_model=sd_model)
         if hasattr(sd_model, 'decoder_pipe') and hasattr(sd_model, 'decoder'):

modules/sd_vae_taesd.py

@@ -36,7 +36,7 @@ prev_cls = ''
 prev_type = ''
 prev_model = ''
 lock = threading.Lock()
-supported = ['sd', 'sdxl', 'f1', 'h1', 'hunyuanvideo', 'wanvideo', 'mochivideo']
+supported = ['sd', 'sdxl', 'f1', 'h1', 'lumina2', 'hunyuanvideo', 'wanvideo', 'mochivideo', 'pixartsigma', 'pixartalpha']
 
 
 def warn_once(msg, variant=None):
@@ -53,11 +53,13 @@ def get_model(model_type = 'decoder', variant = None):
     global prev_cls, prev_type, prev_model # pylint: disable=global-statement
     from modules import shared
     cls = shared.sd_model_type
-    if cls == 'ldm': # original backend
+    if cls in {'ldm', 'pixartalpha'}:
         cls = 'sd'
-    if cls == 'h1': # hidream uses flux vae
+    elif cls in {'h1', 'lumina2'}:
         cls = 'f1'
-    if cls not in supported:
+    elif cls == 'pixartsigma':
+        cls = 'sdxl'
+    elif cls not in supported:
         warn_once(f'cls={shared.sd_model.__class__.__name__} type={cls} unsuppported', variant=variant)
     variant = variant or shared.opts.taesd_variant
     folder = os.path.join(paths.models_path, "TAESD")

modules/sdnq/__init__.py

@@ -0,0 +1,443 @@
+# pylint: disable=redefined-builtin,no-member,protected-access
+
+from typing import Any, Dict, List, Tuple, Optional, Union
+from dataclasses import dataclass
+from enum import Enum
+import torch
+from diffusers.quantizers.base import DiffusersQuantizer
+from diffusers.quantizers.quantization_config import QuantizationConfigMixin
+from diffusers.utils import get_module_from_name
+from modules import devices, shared
+
+from .common import dtype_dict, use_tensorwise_fp8_matmul, quantized_matmul_dtypes, allowed_types, conv_types, conv_transpose_types
+from .dequantizer import dequantizer_dict
+from .forward import get_forward_func
+
+
+def sdnq_quantize_layer(layer, weights_dtype="int8", torch_dtype=None, group_size=0, quant_conv=False, use_quantized_matmul=False, use_quantized_matmul_conv=False, dequantize_fp32=False, quantization_device=None, return_device=None, param_name=None):
+    layer_class_name = layer.__class__.__name__
+    if layer_class_name in allowed_types:
+        is_conv_type = False
+        is_conv_transpose_type = False
+        is_linear_type = False
+        result_shape = None
+        if torch_dtype is None:
+            torch_dtype = devices.dtype
+
+        if layer_class_name in conv_types:
+            if not quant_conv:
+                return layer
+            if dtype_dict[weights_dtype]["num_bits"] < 4:
+                weights_dtype = "uint4"
+            is_conv_type = True
+            reduction_axes = 1
+            output_channel_size, channel_size = layer.weight.shape[:2]
+            group_channel_size = channel_size // layer.groups
+            use_quantized_matmul = False
+            if use_quantized_matmul_conv:
+                use_quantized_matmul = weights_dtype in quantized_matmul_dtypes and group_channel_size >= 32 and output_channel_size >= 32
+                if use_quantized_matmul and not dtype_dict[weights_dtype]["is_integer"]:
+                    use_quantized_matmul = output_channel_size % 16 == 0 and group_channel_size % 16 == 0
+                if use_quantized_matmul:
+                    result_shape = layer.weight.shape
+                    layer.weight.data = layer.weight.reshape(output_channel_size, -1)
+        elif layer_class_name in conv_transpose_types:
+            if not quant_conv:
+                return layer
+            if dtype_dict[weights_dtype]["num_bits"] < 4:
+                weights_dtype = "uint4"
+            is_conv_transpose_type = True
+            reduction_axes = 0
+            channel_size, output_channel_size = layer.weight.shape[:2]
+            use_quantized_matmul = False
+        else:
+            is_linear_type = True
+            reduction_axes = -1
+            output_channel_size, channel_size = layer.weight.shape
+            if use_quantized_matmul:
+                use_quantized_matmul = weights_dtype in quantized_matmul_dtypes and channel_size >= 32 and output_channel_size >= 32
+                if use_quantized_matmul and not dtype_dict[weights_dtype]["is_integer"]:
+                    use_quantized_matmul = output_channel_size % 16 == 0 and channel_size % 16 == 0
+
+        if group_size == 0:
+            if is_linear_type:
+                group_size = 2 ** (2 + dtype_dict[weights_dtype]["num_bits"])
+            else:
+                group_size = 2 ** (1 + dtype_dict[weights_dtype]["num_bits"])
+        elif group_size != -1 and not is_linear_type:
+            group_size = max(group_size // 2, 1)
+
+        if not use_quantized_matmul and group_size > 0:
+            if group_size >= channel_size:
+                group_size = channel_size
+                num_of_groups = 1
+            else:
+                num_of_groups = channel_size // group_size
+                while channel_size % group_size != 0: # find something divisible
+                    num_of_groups -= 1
+                    if num_of_groups <= 1:
+                        group_size = channel_size
+                        num_of_groups = 1
+                        break
+                    group_size = channel_size / num_of_groups
+            group_size = int(group_size)
+            num_of_groups = int(num_of_groups)
+
+            if num_of_groups > 1:
+                result_shape = layer.weight.shape
+                new_shape = list(result_shape)
+                if is_conv_type:
+                    # output_channel_size, channel_size, X, X
+                    # output_channel_size, num_of_groups, group_size, X, X
+                    new_shape[1] = group_size
+                    new_shape.insert(1, num_of_groups)
+                    reduction_axes = 2
+                elif is_conv_transpose_type:
+                    #channel_size, output_channel_size, X, X
+                    #num_of_groups, group_size, output_channel_size, X, X
+                    new_shape[0] = group_size
+                    new_shape.insert(0, num_of_groups)
+                    reduction_axes = 1
+                elif is_linear_type:
+                    # output_channel_size, channel_size
+                    # output_channel_size, num_of_groups, group_size
+                    last_dim_index = layer.weight.ndim
+                    new_shape[last_dim_index - 1 : last_dim_index] = (num_of_groups, group_size)
+                layer.weight.data = layer.weight.reshape(new_shape)
+
+        layer.weight.requires_grad = False
+        if return_device is None:
+            return_device = layer.weight.device
+        if quantization_device is not None:
+            layer.weight.data = layer.weight.to(quantization_device)
+        if layer.weight.dtype != torch.float32:
+            layer.weight.data = layer.weight.to(dtype=torch.float32)
+
+        layer.weight.data, scale, zero_point = quantize_weight(layer.weight, reduction_axes, weights_dtype)
+        if not dequantize_fp32 and not (use_quantized_matmul and not dtype_dict[weights_dtype]["is_integer"] and not use_tensorwise_fp8_matmul):
+            scale = scale.to(torch_dtype)
+            if zero_point is not None:
+                zero_point = zero_point.to(torch_dtype)
+
+        if use_quantized_matmul:
+            scale = scale.transpose(0,1)
+            if dtype_dict[weights_dtype]["num_bits"] == 8:
+                layer.weight.data = layer.weight.transpose(0,1)
+            if not dtype_dict[weights_dtype]["is_integer"]:
+                stride = layer.weight.stride()
+                if stride[0] > stride[1] and stride[1] == 1:
+                    layer.weight.data = layer.weight.t().contiguous().t()
+                if not use_tensorwise_fp8_matmul:
+                    scale = scale.to(torch.float32)
+
+        layer.sdnq_dequantizer = dequantizer_dict[weights_dtype](
+            scale=scale,
+            zero_point=zero_point,
+            quantized_weight_shape=layer.weight.shape,
+            result_dtype=torch_dtype,
+            result_shape=result_shape,
+            weights_dtype=weights_dtype,
+            use_quantized_matmul=use_quantized_matmul,
+        )
+        layer.weight.data = layer.sdnq_dequantizer.pack_weight(layer.weight).to(return_device)
+        layer.sdnq_dequantizer = layer.sdnq_dequantizer.to(return_device)
+
+        layer.forward = get_forward_func(layer_class_name, use_quantized_matmul, dtype_dict[weights_dtype]["is_integer"], use_tensorwise_fp8_matmul)
+        layer.forward = layer.forward.__get__(layer, layer.__class__)
+        devices.torch_gc(force=False, reason=f"SDNQ param_name: {param_name}")
+    return layer
+
+
+def apply_sdnq_to_module(model, weights_dtype="int8", torch_dtype=None, group_size=0, quant_conv=False, use_quantized_matmul=False, use_quantized_matmul_conv=False, dequantize_fp32=False, quantization_device=None, return_device=None, param_name=None): # pylint: disable=unused-argument
+    has_children = list(model.children())
+    if not has_children:
+        return model
+    for module_param_name, module in model.named_children():
+        if hasattr(module, "weight") and module.weight is not None:
+            module = sdnq_quantize_layer(
+                module,
+                weights_dtype=weights_dtype,
+                torch_dtype=torch_dtype,
+                group_size=group_size,
+                quant_conv=quant_conv,
+                use_quantized_matmul=use_quantized_matmul,
+                use_quantized_matmul_conv=use_quantized_matmul_conv,
+                dequantize_fp32=dequantize_fp32,
+                quantization_device=quantization_device,
+                return_device=return_device,
+                param_name=module_param_name,
+            )
+        module = apply_sdnq_to_module(
+                module,
+                weights_dtype=weights_dtype,
+                torch_dtype=torch_dtype,
+                group_size=group_size,
+                quant_conv=quant_conv,
+                use_quantized_matmul=use_quantized_matmul,
+                use_quantized_matmul_conv=use_quantized_matmul_conv,
+                dequantize_fp32=dequantize_fp32,
+                quantization_device=quantization_device,
+                return_device=return_device,
+                param_name=module_param_name,
+            )
+    return model
+
+
+def get_scale_asymmetric(weight: torch.FloatTensor, reduction_axes: Union[int, List[int]], weights_dtype: str) -> Tuple[torch.FloatTensor, torch.FloatTensor]:
+    zero_point = torch.amin(weight, dim=reduction_axes, keepdims=True)
+    scale = torch.amax(weight, dim=reduction_axes, keepdims=True).sub_(zero_point).div_(dtype_dict[weights_dtype]["max"] - dtype_dict[weights_dtype]["min"])
+    eps = torch.finfo(scale.dtype).eps # prevent divison by 0
+    scale = torch.where(torch.abs(scale) < eps, eps, scale)
+    if dtype_dict[weights_dtype]["min"] != 0:
+        zero_point.sub_(torch.mul(scale, dtype_dict[weights_dtype]["min"]))
+    return scale, zero_point
+
+
+def get_scale_symmetric(weight: torch.FloatTensor, reduction_axes: Union[int, List[int]], weights_dtype: str) -> torch.FloatTensor:
+    scale = torch.amax(weight.abs(), dim=reduction_axes, keepdims=True).div_(dtype_dict[weights_dtype]["max"])
+    eps = torch.finfo(scale.dtype).eps # prevent divison by 0
+    scale = torch.where(torch.abs(scale) < eps, eps, scale)
+    return scale
+
+
+def quantize_weight(weight: torch.FloatTensor, reduction_axes: Union[int, List[int]], weights_dtype: str):
+    if dtype_dict[weights_dtype]["is_unsigned"]:
+        scale, zero_point = get_scale_asymmetric(weight, reduction_axes, weights_dtype)
+        quantized_weight = torch.sub(weight, zero_point).div_(scale)
+    else:
+        scale = get_scale_symmetric(weight, reduction_axes, weights_dtype)
+        quantized_weight = torch.div(weight, scale)
+        zero_point = None
+    if dtype_dict[weights_dtype]["is_integer"]:
+        quantized_weight.round_()
+    quantized_weight = quantized_weight.clamp_(dtype_dict[weights_dtype]["min"], dtype_dict[weights_dtype]["max"]).to(dtype_dict[weights_dtype]["torch_dtype"])
+    return quantized_weight, scale, zero_point
+
+
+class QuantizationMethod(str, Enum):
+    SDNQ = "sdnq"
+
+
+class SDNQQuantizer(DiffusersQuantizer):
+    r"""
+    Diffusers Quantizer for SDNQ
+    """
+
+    requires_parameters_quantization = True
+    use_keep_in_fp32_modules = True
+    requires_calibration = False
+    required_packages = None
+    torch_dtype = None
+
+    def __init__(self, quantization_config, **kwargs): # pylint: disable=useless-parent-delegation
+        super().__init__(quantization_config, **kwargs)
+        self.modules_to_not_convert = []
+
+    def check_if_quantized_param(
+        self,
+        model,
+        param_value: "torch.Tensor",
+        param_name: str,
+        state_dict: Dict[str, Any], # pylint: disable=unused-argument
+        **kwargs, # pylint: disable=unused-argument
+    ):
+        if param_name.endswith(".weight"):
+            split_param_name = param_name.split(".")
+            if param_name not in self.modules_to_not_convert and not any(param in split_param_name for param in self.modules_to_not_convert):
+                layer_class_name = get_module_from_name(model, param_name)[0].__class__.__name__
+                if layer_class_name in allowed_types:
+                    if layer_class_name in conv_types or layer_class_name in conv_transpose_types:
+                        if self.quantization_config.quant_conv:
+                            return True
+                    else:
+                        return True
+        param_value.data = param_value.clone() # safetensors is unable to release the cpu memory without this
+        return False
+
+    def check_quantized_param(self, *args, **kwargs) -> bool:
+        """
+        needed for transformers compatibilty, returns self.check_if_quantized_param
+        """
+        return self.check_if_quantized_param(*args, **kwargs)
+
+    def create_quantized_param( # pylint: disable=arguments-differ
+        self,
+        model,
+        param_value: torch.FloatTensor,
+        param_name: str,
+        target_device: torch.device,
+        state_dict: Dict[str, Any], # pylint: disable=unused-argument
+        unexpected_keys: List[str], # pylint: disable=unused-argument
+        **kwargs, # pylint: disable=unused-argument
+    ):
+        if self.quantization_config.return_device is not None:
+            return_device = self.quantization_config.return_device
+        else:
+            return_device = target_device
+
+        if self.quantization_config.quantization_device is not None:
+            target_device = self.quantization_config.quantization_device
+
+        if param_value.dtype == torch.float32 and devices.same_device(param_value.device, target_device):
+            param_value = param_value.clone()
+        else:
+            param_value = param_value.to(target_device).to(dtype=torch.float32)
+
+        layer, _ = get_module_from_name(model, param_name)
+        layer.weight = torch.nn.Parameter(param_value, requires_grad=False)
+        layer = sdnq_quantize_layer(
+            layer,
+            weights_dtype=self.quantization_config.weights_dtype,
+            torch_dtype=self.torch_dtype,
+            group_size=self.quantization_config.group_size,
+            quant_conv=self.quantization_config.quant_conv,
+            use_quantized_matmul=self.quantization_config.use_quantized_matmul,
+            use_quantized_matmul_conv=self.quantization_config.use_quantized_matmul_conv,
+            dequantize_fp32=self.quantization_config.dequantize_fp32,
+            quantization_device=None,
+            return_device=return_device,
+            param_name=param_name,
+        )
+
+    def adjust_max_memory(self, max_memory: Dict[str, Union[int, str]]) -> Dict[str, Union[int, str]]:
+        max_memory = {key: val * 0.80 for key, val in max_memory.items()}
+        return max_memory
+
+    def adjust_target_dtype(self, target_dtype: torch.dtype) -> torch.dtype: # pylint: disable=unused-argument,arguments-renamed
+        return dtype_dict[self.quantization_config.weights_dtype]["target_dtype"]
+
+    def update_torch_dtype(self, torch_dtype: torch.dtype = None) -> torch.dtype:
+        if torch_dtype is None:
+            torch_dtype = devices.dtype
+        self.torch_dtype = torch_dtype
+        return torch_dtype
+
+    def _process_model_before_weight_loading( # pylint: disable=arguments-differ
+        self,
+        model,
+        device_map, # pylint: disable=unused-argument
+        keep_in_fp32_modules: List[str] = [],
+        **kwargs, # pylint: disable=unused-argument
+    ):
+        if keep_in_fp32_modules is not None:
+            self.modules_to_not_convert.extend(keep_in_fp32_modules)
+        self.modules_to_not_convert.extend(self.quantization_config.modules_to_not_convert)
+        self.quantization_config.modules_to_not_convert = self.modules_to_not_convert
+        model.config.quantization_config = self.quantization_config
+
+    def _process_model_after_weight_loading(self, model, **kwargs): # pylint: disable=unused-argument
+        if shared.opts.diffusers_offload_mode != "none":
+            model = model.to(devices.cpu)
+        devices.torch_gc(force=True)
+        return model
+
+    def get_cuda_warm_up_factor(self):
+        return 32 // dtype_dict[self.quantization_config.weights_dtype]["num_bits"]
+
+    def update_tp_plan(self, config):
+        """
+        needed for transformers compatibilty, no-op function
+        """
+        return config
+
+    def update_unexpected_keys(self, model, unexpected_keys: List[str], prefix: str) -> List[str]: # pylint: disable=unused-argument
+        """
+        needed for transformers compatibilty, no-op function
+        """
+        return unexpected_keys
+
+    def update_missing_keys_after_loading(self, model, missing_keys: List[str], prefix: str) -> List[str]: # pylint: disable=unused-argument
+        """
+        needed for transformers compatibilty, no-op function
+        """
+        return missing_keys
+
+    def update_expected_keys(self, model, expected_keys: List[str], loaded_keys: List[str]) -> List[str]: # pylint: disable=unused-argument
+        """
+        needed for transformers compatibilty, no-op function
+        """
+        return expected_keys
+
+    @property
+    def is_trainable(self):
+        return False
+
+    @property
+    def is_serializable(self):
+        return True
+
+    @property
+    def is_compileable(self):
+        return True
+
+
+@dataclass
+class SDNQConfig(QuantizationConfigMixin):
+    """
+    This is a wrapper class about all possible attributes and features that you can play with a model that has been
+    loaded using `sdnq`.
+
+    Args:
+        weights_dtype (`str`, *optional*, defaults to `"int8"`):
+            The target dtype for the weights after quantization. Supported values are:
+            ("int8", "int7", "int6", "int5", "int4", "int3", "int2", "uint8", "uint7", "uint6", "uint5", "uint4", "uint3", "uint2", "uint1", "bool", "float8_e4m3fn", "float8_e4m3fnuz", "float8_e5m2", "float8_e5m2fnuz")
+       modules_to_not_convert (`list`, *optional*, default to `None`):
+            The list of modules to not quantize, useful for quantizing models that explicitly require to have some
+            modules left in their original precision (e.g. Whisper encoder, Llava encoder, Mixtral gate layers).
+    """
+
+    def __init__( # pylint: disable=super-init-not-called
+        self,
+        weights_dtype: str = "int8",
+        group_size: int = 0,
+        quant_conv: bool = False,
+        use_quantized_matmul: bool = False,
+        use_quantized_matmul_conv: bool = False,
+        dequantize_fp32: bool = False,
+        quantization_device: Optional[torch.device] = None,
+        return_device: Optional[torch.device] = None,
+        modules_to_not_convert: Optional[List[str]] = None,
+        **kwargs, # pylint: disable=unused-argument
+    ):
+        self.weights_dtype = weights_dtype
+        self.quant_method = QuantizationMethod.SDNQ
+        self.group_size = group_size
+        self.quant_conv = quant_conv
+        self.use_quantized_matmul = use_quantized_matmul
+        self.use_quantized_matmul_conv = use_quantized_matmul_conv
+        self.dequantize_fp32 = dequantize_fp32
+        self.quantization_device = quantization_device
+        self.return_device = return_device
+        self.modules_to_not_convert = modules_to_not_convert
+        self.post_init()
+        self.is_integer = dtype_dict[self.weights_dtype]["is_integer"]
+
+    def post_init(self):
+        r"""
+        Safety checker that arguments are correct
+        """
+        accepted_weights = ["int8", "int7", "int6", "int5", "int4", "int3", "int2", "uint8", "uint7", "uint6", "uint5", "uint4", "uint3", "uint2", "uint1", "bool", "float8_e4m3fn", "float8_e4m3fnuz", "float8_e5m2", "float8_e5m2fnuz"]
+        if self.weights_dtype not in accepted_weights:
+            raise ValueError(f"Only support weights in {accepted_weights} but found {self.weights_dtype}")
+        if not isinstance(self.modules_to_not_convert, list):
+            self.modules_to_not_convert = [self.modules_to_not_convert]
+
+
+class SDNQ_T5DenseGatedActDense(torch.nn.Module): # forward can't find what self is without creating a class
+    def __init__(self, T5DenseGatedActDense, dtype):
+        super().__init__()
+        self.wi_0 = T5DenseGatedActDense.wi_0
+        self.wi_1 = T5DenseGatedActDense.wi_1
+        self.wo = T5DenseGatedActDense.wo
+        self.dropout = T5DenseGatedActDense.dropout
+        self.act = T5DenseGatedActDense.act
+        self.torch_dtype = dtype
+
+    def forward(self, hidden_states):
+        hidden_gelu = self.act(self.wi_0(hidden_states))
+        hidden_linear = self.wi_1(hidden_states)
+        hidden_states = hidden_gelu * hidden_linear
+        hidden_states = self.dropout(hidden_states)
+        hidden_states = hidden_states.to(self.torch_dtype) # this line needs to be forced
+        hidden_states = self.wo(hidden_states)
+        return hidden_states

modules/sdnq/common.py

@@ -0,0 +1,41 @@
+# pylint: disable=redefined-builtin,no-member,protected-access
+
+import sys
+import torch
+from accelerate.utils import CustomDtype
+from modules import devices
+
+torch_version = float(torch.__version__[:3])
+
+dtype_dict = {
+    "int8": {"min": -128, "max": 127, "num_bits": 8, "target_dtype": torch.int8, "torch_dtype": torch.int8, "storage_dtype": torch.int8, "is_unsigned": False, "is_integer": True},
+    "int7": {"min": -64, "max": 63, "num_bits": 7, "target_dtype": torch.int8, "torch_dtype": torch.int8, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": True},
+    "int6": {"min": -32, "max": 31, "num_bits": 6, "target_dtype": torch.int8, "torch_dtype": torch.int8, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": True},
+    "int5": {"min": -16, "max": 15, "num_bits": 5, "target_dtype": CustomDtype.INT4, "torch_dtype": torch.int8, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": True},
+    "int4": {"min": -8, "max": 7, "num_bits": 4, "target_dtype": CustomDtype.INT4, "torch_dtype": torch.int8, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": True},
+    "int3": {"min": -4, "max": 3, "num_bits": 3, "target_dtype": CustomDtype.INT4, "torch_dtype": torch.int8, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": True},
+    "int2": {"min": -2, "max": 1, "num_bits": 2, "target_dtype": CustomDtype.INT2, "torch_dtype": torch.int8, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": True},
+    "uint8": {"min": 0, "max": 255, "num_bits": 8, "target_dtype": torch.uint8, "torch_dtype": torch.uint8, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": True},
+    "uint7": {"min": 0, "max": 127, "num_bits": 7, "target_dtype": torch.uint8, "torch_dtype": torch.uint8, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": True},
+    "uint6": {"min": 0, "max": 63, "num_bits": 6, "target_dtype": torch.uint8, "torch_dtype": torch.uint8, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": True},
+    "uint5": {"min": 0, "max": 31, "num_bits": 5, "target_dtype": CustomDtype.INT4, "torch_dtype": torch.uint8, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": True},
+    "uint4": {"min": 0, "max": 15, "num_bits": 4, "target_dtype": CustomDtype.INT4, "torch_dtype": torch.uint8, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": True},
+    "uint3": {"min": 0, "max": 7, "num_bits": 3, "target_dtype": CustomDtype.INT4, "torch_dtype": torch.uint8, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": True},
+    "uint2": {"min": 0, "max": 3, "num_bits": 2, "target_dtype": CustomDtype.INT2, "torch_dtype": torch.uint8, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": True},
+    "uint1": {"min": 0, "max": 1, "num_bits": 1, "target_dtype": torch.bool, "torch_dtype": torch.bool, "storage_dtype": torch.bool, "is_unsigned": True, "is_integer": True},
+    "float8_e4m3fn": {"min": -448, "max": 448, "num_bits": 8, "target_dtype": torch.float8_e4m3fn, "torch_dtype": torch.float8_e4m3fn, "storage_dtype": torch.float8_e4m3fn, "is_unsigned": False, "is_integer": False},
+    "float8_e5m2": {"min": -57344, "max": 57344, "num_bits": 8, "target_dtype": torch.float8_e5m2, "torch_dtype": torch.float8_e5m2, "storage_dtype": torch.float8_e5m2, "is_unsigned": False, "is_integer": False},
+    "float8_e4m3fnuz": {"min": -240, "max": 240, "num_bits": 8, "target_dtype": CustomDtype.FP8, "torch_dtype": torch.float8_e4m3fnuz, "storage_dtype": torch.float8_e4m3fnuz, "is_unsigned": False, "is_integer": False},
+    "float8_e5m2fnuz": {"min": -57344, "max": 57344, "num_bits": 8, "target_dtype": CustomDtype.FP8, "torch_dtype": torch.float8_e5m2fnuz, "storage_dtype": torch.float8_e5m2fnuz, "is_unsigned": False, "is_integer": False},
+}
+dtype_dict["bool"] = dtype_dict["uint1"]
+
+use_tensorwise_fp8_matmul = torch_version < 2.5 or devices.backend in {"cpu", "openvino"} or (devices.backend == "cuda" and sys.platform == "win32" and torch_version <= 2.7 and torch.cuda.get_device_capability(devices.device) == (8,9))
+quantized_matmul_dtypes = ("int8", "int7", "int6", "int5", "int4", "int3", "int2", "float8_e4m3fn", "float8_e5m2")
+if devices.backend in {"cpu", "openvino"}:
+    quantized_matmul_dtypes += ("float8_e4m3fnuz", "float8_e5m2fnuz")
+
+linear_types = ("Linear",)
+conv_types = ("Conv1d", "Conv2d", "Conv3d")
+conv_transpose_types = ("ConvTranspose1d", "ConvTranspose2d", "ConvTranspose3d")
+allowed_types = linear_types + conv_types + conv_transpose_types

modules/sdnq/dequantizer.py

@@ -0,0 +1,185 @@
+# pylint: disable=redefined-builtin,no-member,protected-access
+
+import torch
+from modules import shared
+
+from .common import dtype_dict
+from .packed_int import pack_int_symetric, unpack_int_symetric, packed_int_function_dict
+
+
+def dequantize_asymmetric(input: torch.ByteTensor, scale: torch.FloatTensor, zero_point: torch.FloatTensor, dtype: torch.dtype, result_shape: torch.Size) -> torch.FloatTensor:
+    result = torch.addcmul(zero_point, input.to(dtype=scale.dtype), scale).to(dtype=dtype)
+    if result_shape is not None:
+        result = result.reshape(result_shape)
+    return result
+
+
+def dequantize_symmetric(input: torch.CharTensor, scale: torch.FloatTensor, dtype: torch.dtype, result_shape: torch.Size, skip_quantized_matmul: bool = False) -> torch.FloatTensor:
+    if skip_quantized_matmul:
+        result = input.transpose(0,1).to(dtype=scale.dtype).mul_(scale.transpose(0,1)).to(dtype=dtype)
+    else:
+        result = input.to(dtype=scale.dtype).mul_(scale).to(dtype=dtype)
+    if result_shape is not None:
+        result = result.reshape(result_shape)
+    return result
+
+
+def dequantize_symmetric_with_bias(input: torch.CharTensor, bias: torch.FloatTensor, scale: torch.FloatTensor, dtype: torch.dtype, result_shape: torch.Size) -> torch.FloatTensor:
+    return torch.addcmul(bias, input.to(dtype=scale.dtype), scale).to(dtype=dtype).reshape(result_shape)
+
+
+def dequantize_packed_int_asymmetric(input: torch.ByteTensor, scale: torch.FloatTensor, zero_point: torch.FloatTensor, shape: torch.Size, dtype: torch.dtype, result_shape: torch.Size, weights_dtype: str) -> torch.FloatTensor:
+    return dequantize_asymmetric(packed_int_function_dict[weights_dtype]["unpack"](input, shape), scale, zero_point, dtype, result_shape)
+
+
+def dequantize_packed_int_symmetric(input: torch.ByteTensor, scale: torch.FloatTensor, shape: torch.Size, dtype: torch.dtype, result_shape: torch.Size, weights_dtype: str, skip_quantized_matmul: bool = False) -> torch.FloatTensor:
+    if skip_quantized_matmul:
+        return dequantize_symmetric(unpack_int_symetric(input, shape, weights_dtype, dtype=scale.dtype), scale.transpose(0,1), dtype, result_shape)
+    else:
+        return dequantize_symmetric(unpack_int_symetric(input, shape, weights_dtype, dtype=scale.dtype), scale, dtype, result_shape)
+
+
+class AsymmetricWeightsDequantizer(torch.nn.Module):
+    def __init__(
+        self,
+        scale: torch.Tensor,
+        zero_point: torch.Tensor,
+        result_dtype: torch.dtype,
+        result_shape: torch.Size,
+        weights_dtype: str,
+        **kwargs, # pylint: disable=unused-argument
+    ):
+        super().__init__()
+        self.weights_dtype = weights_dtype
+        self.use_quantized_matmul = False
+        self.result_dtype = result_dtype
+        self.result_shape = result_shape
+        self.register_buffer("scale", scale)
+        self.register_buffer("zero_point", zero_point)
+
+    def pack_weight(self, weight: torch.Tensor) -> torch.Tensor:
+        return weight.to(dtype=dtype_dict[self.weights_dtype]["torch_dtype"])
+
+    def forward(self, weight, **kwargs): # pylint: disable=unused-argument
+        return dequantize_asymmetric_compiled(weight, self.scale, self.zero_point, self.result_dtype, self.result_shape)
+
+
+class SymmetricWeightsDequantizer(torch.nn.Module):
+    def __init__(
+        self,
+        scale: torch.Tensor,
+        result_dtype: torch.dtype,
+        result_shape: torch.Size,
+        weights_dtype: str,
+        use_quantized_matmul: bool = False,
+        **kwargs, # pylint: disable=unused-argument
+    ):
+        super().__init__()
+        self.weights_dtype = weights_dtype
+        self.use_quantized_matmul = use_quantized_matmul
+        self.result_dtype = result_dtype
+        self.result_shape = result_shape
+        self.register_buffer("scale", scale)
+
+    def pack_weight(self, weight: torch.Tensor) -> torch.Tensor:
+        return weight.to(dtype=dtype_dict[self.weights_dtype]["torch_dtype"])
+
+    def forward(self, weight, skip_quantized_matmul=False, **kwargs): # pylint: disable=unused-argument
+        return dequantize_symmetric_compiled(weight, self.scale, self.result_dtype, self.result_shape, skip_quantized_matmul=skip_quantized_matmul)
+
+
+class PackedINTAsymmetricWeightsDequantizer(torch.nn.Module):
+    def __init__(
+        self,
+        scale: torch.Tensor,
+        zero_point: torch.Tensor,
+        quantized_weight_shape: torch.Size,
+        result_dtype: torch.dtype,
+        result_shape: torch.Size,
+        weights_dtype: str,
+        **kwargs, # pylint: disable=unused-argument
+    ):
+        super().__init__()
+        self.weights_dtype = weights_dtype
+        self.use_quantized_matmul = False
+        self.quantized_weight_shape = quantized_weight_shape
+        self.result_dtype = result_dtype
+        self.result_shape = result_shape
+        self.register_buffer("scale", scale)
+        self.register_buffer("zero_point", zero_point)
+
+    def pack_weight(self, weight: torch.Tensor) -> torch.Tensor:
+        return packed_int_function_dict[self.weights_dtype]["pack"](weight.to(dtype=dtype_dict[self.weights_dtype]["torch_dtype"]))
+
+    def forward(self, weight, **kwargs): # pylint: disable=unused-argument
+        return dequantize_packed_int_asymmetric_compiled(weight, self.scale, self.zero_point, self.quantized_weight_shape, self.result_dtype, self.result_shape, self.weights_dtype)
+
+
+class PackedINTSymmetricWeightsDequantizer(torch.nn.Module):
+    def __init__(
+        self,
+        scale: torch.Tensor,
+        quantized_weight_shape: torch.Size,
+        result_dtype: torch.dtype,
+        result_shape: torch.Size,
+        weights_dtype: str,
+        use_quantized_matmul: bool = False,
+        **kwargs, # pylint: disable=unused-argument
+    ):
+        super().__init__()
+        self.weights_dtype = weights_dtype
+        self.use_quantized_matmul = use_quantized_matmul
+        self.quantized_weight_shape = quantized_weight_shape
+        self.result_dtype = result_dtype
+        self.result_shape = result_shape
+        self.register_buffer("scale", scale)
+
+    def pack_weight(self, weight: torch.Tensor) -> torch.Tensor:
+        return pack_int_symetric(weight, self.weights_dtype)
+
+    def forward(self, weight, skip_quantized_matmul=False, **kwargs): # pylint: disable=unused-argument
+        return dequantize_packed_int_symmetric_compiled(weight, self.scale, self.quantized_weight_shape, self.result_dtype, self.result_shape, self.weights_dtype, skip_quantized_matmul=skip_quantized_matmul)
+
+
+dequantizer_dict = {
+    "int8": SymmetricWeightsDequantizer,
+    "int7": PackedINTSymmetricWeightsDequantizer,
+    "int6": PackedINTSymmetricWeightsDequantizer,
+    "int5": PackedINTSymmetricWeightsDequantizer,
+    "int4": PackedINTSymmetricWeightsDequantizer,
+    "int3": PackedINTSymmetricWeightsDequantizer,
+    "int2": PackedINTSymmetricWeightsDequantizer,
+    "uint8": AsymmetricWeightsDequantizer,
+    "uint7": PackedINTAsymmetricWeightsDequantizer,
+    "uint6": PackedINTAsymmetricWeightsDequantizer,
+    "uint5": PackedINTAsymmetricWeightsDequantizer,
+    "uint4": PackedINTAsymmetricWeightsDequantizer,
+    "uint3": PackedINTAsymmetricWeightsDequantizer,
+    "uint2": PackedINTAsymmetricWeightsDequantizer,
+    "uint1": AsymmetricWeightsDequantizer,
+    "bool": AsymmetricWeightsDequantizer,
+    "float8_e4m3fn": SymmetricWeightsDequantizer,
+    "float8_e4m3fnuz": SymmetricWeightsDequantizer,
+    "float8_e5m2": SymmetricWeightsDequantizer,
+    "float8_e5m2fnuz": SymmetricWeightsDequantizer,
+}
+
+
+if shared.opts.sdnq_dequantize_compile:
+    try:
+        torch._dynamo.config.cache_size_limit = max(8192, torch._dynamo.config.cache_size_limit)
+        dequantize_asymmetric_compiled = torch.compile(dequantize_asymmetric, fullgraph=True)
+        dequantize_symmetric_compiled = torch.compile(dequantize_symmetric, fullgraph=True)
+        dequantize_packed_int_asymmetric_compiled = torch.compile(dequantize_packed_int_asymmetric, fullgraph=True)
+        dequantize_packed_int_symmetric_compiled = torch.compile(dequantize_packed_int_symmetric, fullgraph=True)
+    except Exception as e:
+        shared.log.warning(f"Quantization: type=sdnq Dequantize using torch.compile is not available: {e}")
+        dequantize_asymmetric_compiled = dequantize_asymmetric
+        dequantize_symmetric_compiled = dequantize_symmetric
+        dequantize_packed_int_asymmetric_compiled = dequantize_packed_int_asymmetric
+        dequantize_packed_int_symmetric_compiled = dequantize_packed_int_symmetric
+else:
+    dequantize_asymmetric_compiled = dequantize_asymmetric
+    dequantize_symmetric_compiled = dequantize_symmetric
+    dequantize_packed_int_asymmetric_compiled = dequantize_packed_int_asymmetric
+    dequantize_packed_int_symmetric_compiled = dequantize_packed_int_symmetric

modules/sdnq/forward.py

@@ -0,0 +1,409 @@
+# pylint: disable=redefined-builtin,no-member,protected-access
+
+from typing import Callable, List, Tuple, Optional
+import torch
+from modules import shared
+
+from .common import conv_types, conv_transpose_types
+from .dequantizer import dequantize_symmetric, dequantize_symmetric_with_bias
+from .packed_int import unpack_int_symetric
+
+
+def get_forward_func(layer_class_name: str, use_quantized_matmul: bool, is_integer: bool, use_tensorwise_fp8_matmul: bool) -> Callable: # pylint: disable=inconsistent-return-statements
+    if layer_class_name in conv_types:
+        if use_quantized_matmul:
+            if is_integer:
+                return quantized_conv_forward_int8_matmul
+            else:
+                if use_tensorwise_fp8_matmul:
+                    return quantized_conv_forward_fp8_matmul_tensorwise
+                else:
+                    return quantized_conv_forward_fp8_matmul
+        else:
+            return quantized_conv_forward
+    elif layer_class_name in conv_transpose_types:
+        if layer_class_name.endswith("1d"):
+            return quantized_conv_transpose_1d_forward
+        elif layer_class_name.endswith("2d"):
+            return quantized_conv_transpose_2d_forward
+        elif layer_class_name.endswith("3d"):
+            return quantized_conv_transpose_3d_forward
+    else:
+        if use_quantized_matmul:
+            if is_integer:
+                return quantized_linear_forward_int8_matmul
+            else:
+                if use_tensorwise_fp8_matmul:
+                    return quantized_linear_forward_fp8_matmul_tensorwise
+                else:
+                    return quantized_linear_forward_fp8_matmul
+        else:
+            return quantized_linear_forward
+
+
+def quantize_fp8_matmul_input(input: torch.FloatTensor) -> Tuple[torch.Tensor, torch.FloatTensor]:
+    input = input.flatten(0,-2).contiguous()
+    input_scale = torch.amax(input.abs(), dim=-1, keepdims=True).div_(448)
+    input = torch.div(input, input_scale).clamp_(-448, 448).to(torch.float8_e4m3fn)
+    input_scale = input_scale.to(torch.float32)
+    return input, input_scale
+
+
+def quantize_fp8_matmul_input_tensorwise(input: torch.FloatTensor, scale: torch.FloatTensor) -> Tuple[torch.Tensor, torch.FloatTensor]:
+    input = input.flatten(0,-2).contiguous()
+    input_scale = torch.amax(input.abs(), dim=-1, keepdims=True).div_(448)
+    input = torch.div(input, input_scale).clamp_(-448, 448).to(torch.float8_e4m3fn)
+    scale = torch.mul(input_scale, scale)
+    if scale.dtype == torch.float16: # fp16 will overflow
+        scale = scale.to(dtype=torch.float32)
+    return input, scale
+
+
+def quantize_int8_matmul_input(input: torch.FloatTensor, scale: torch.FloatTensor) -> Tuple[torch.CharTensor, torch.FloatTensor]:
+    input = input.flatten(0,-2).contiguous()
+    input_scale = torch.amax(input.abs(), dim=-1, keepdims=True).div_(127)
+    input = torch.div(input, input_scale).round_().clamp_(-128, 127).to(torch.int8)
+    scale = torch.mul(input_scale, scale)
+    if scale.dtype == torch.float16: # fp16 will overflow
+        scale = scale.to(dtype=torch.float32)
+    return input, scale
+
+
+def fp8_matmul(
+    input: torch.FloatTensor,
+    weight: torch.Tensor,
+    bias: torch.FloatTensor,
+    scale: torch.FloatTensor,
+) -> torch.FloatTensor:
+    return_dtype = input.dtype
+    output_shape = list(input.shape)
+    output_shape[-1] = weight.shape[-1]
+    input, input_scale = quantize_fp8_matmul_input(input)
+    return torch._scaled_mm(input, weight, scale_a=input_scale, scale_b=scale, bias=bias, out_dtype=return_dtype).reshape(output_shape)
+
+
+# sm89 doesn't support row wise scale in Windows
+def fp8_matmul_tensorwise(
+    input: torch.FloatTensor,
+    weight: torch.Tensor,
+    bias: torch.FloatTensor,
+    scale: torch.FloatTensor,
+) -> torch.FloatTensor:
+    return_dtype = input.dtype
+    output_shape = list(input.shape)
+    output_shape[-1] = weight.shape[-1]
+    dummy_input_scale = torch.ones(1, device=input.device, dtype=torch.float32)
+    input, scale = quantize_fp8_matmul_input_tensorwise(input, scale)
+    if bias is not None:
+        return dequantize_symmetric_with_bias(torch._scaled_mm(input, weight, scale_a=dummy_input_scale, scale_b=dummy_input_scale, bias=None, out_dtype=scale.dtype), bias, scale, return_dtype, output_shape)
+    else:
+        return dequantize_symmetric(torch._scaled_mm(input, weight, scale_a=dummy_input_scale, scale_b=dummy_input_scale, bias=None, out_dtype=scale.dtype), scale, return_dtype, output_shape)
+
+
+def int8_matmul(
+    input: torch.FloatTensor,
+    weight: torch.Tensor,
+    bias: torch.FloatTensor,
+    scale: torch.FloatTensor,
+    quantized_weight_shape: torch.Size,
+    weights_dtype: str,
+) -> torch.FloatTensor:
+    if quantized_weight_shape is not None:
+        weight = unpack_int_symetric(weight, quantized_weight_shape, weights_dtype, dtype=torch.int8, transpose=True)
+    return_dtype = input.dtype
+    output_shape = list(input.shape)
+    output_shape[-1] = weight.shape[-1]
+    input, scale = quantize_int8_matmul_input(input, scale)
+    if bias is not None:
+        return dequantize_symmetric_with_bias(torch._int_mm(input, weight), bias, scale, return_dtype, output_shape)
+    else:
+        return dequantize_symmetric(torch._int_mm(input, weight), scale, return_dtype, output_shape)
+
+
+def process_conv_input(conv_type, input, reversed_padding_repeated_twice, padding_mode, result_shape, stride, padding, dilation):
+    if conv_type == 1:
+        batch_size, _, L_in = input.shape
+        C_out, _, K_l = result_shape
+        L_out = (L_in + 2 * padding[1] - dilation[1] * (K_l - 1) - 1) // stride[1] + 1
+        mm_output_shape = (batch_size, L_out, C_out)
+        kernel_size = (1, K_l)
+    if conv_type == 2:
+        batch_size, _, H_in, W_in = input.shape
+        C_out, _, K_h, K_w = result_shape
+        H_out = (H_in + 2 * padding[0] - dilation[0] * (K_h - 1) - 1) // stride[0] + 1
+        W_out = (W_in + 2 * padding[1] - dilation[1] * (K_w - 1) - 1) // stride[1] + 1
+        mm_output_shape = (batch_size, H_out, W_out, C_out)
+        kernel_size = (K_h, K_w)
+    else:
+        batch_size, _, D_in, H_in, W_in = input.shape
+        C_out, _, K_d, K_h, K_w = result_shape
+        D_out = (D_in + 2 * padding[0] - dilation[0] * (K_d - 1) - 1) // stride[0] + 1
+        H_out = (H_in + 2 * padding[1] - dilation[1] * (K_h - 1) - 1) // stride[1] + 1
+        W_out = (W_in + 2 * padding[2] - dilation[2] * (K_w - 1) - 1) // stride[2] + 1
+        mm_output_shape = (batch_size, D_out, H_out, W_out, C_out)
+        kernel_size = (K_d, K_h, K_w)
+
+    if padding_mode != "zeros":
+        input = torch.nn.functional.pad(input, reversed_padding_repeated_twice, mode=padding_mode)
+        padding = (0,) * (conv_type if conv_type != 1 else 2)
+    elif conv_type == 3:
+        input = torch.nn.functional.pad(input, reversed_padding_repeated_twice)
+
+    if conv_type == 1:
+        input = input.unsqueeze(2)
+
+    if conv_type == 3:
+        K_D_eff = K_d + (K_d - 1) * (dilation[0] - 1)
+        K_H_eff = K_h + (K_h - 1) * (dilation[0] - 1)
+        K_W_eff = K_w + (K_w - 1) * (dilation[0] - 1)
+        input = input.unfold(2, K_D_eff, stride[0]).unfold(3, K_H_eff, stride[1]).unfold(4, K_W_eff, stride[2])
+        if dilation[0] > 1:
+            input = input[..., ::dilation[0], :, :]
+        if dilation[1] > 1:
+            input = input[..., ::dilation[1], :]
+        if dilation[2] > 1:
+            input = input[..., ::dilation[2]]
+        input = input.permute(0, 2, 3, 4, 1, 5, 6, 7).reshape(batch_size, D_out * H_out * W_out, -1)
+    else:
+        input = torch.nn.functional.unfold(input, kernel_size=kernel_size, padding=padding, stride=stride, dilation=dilation).transpose(1,2)
+    return input, mm_output_shape
+
+
+def conv_fp8_matmul(
+    input: torch.FloatTensor,
+    weight: torch.Tensor,
+    bias: torch.FloatTensor,
+    scale: torch.FloatTensor,
+    result_shape: torch.Size,
+    reversed_padding_repeated_twice: List[int],
+    padding_mode: str, conv_type: int,
+    groups: int, stride: List[int],
+    padding: List[int], dilation: List[int],
+) -> torch.FloatTensor:
+    return_dtype = input.dtype
+    input, mm_output_shape = process_conv_input(conv_type, input, reversed_padding_repeated_twice, padding_mode, result_shape, stride, padding, dilation)
+    input, input_scale = quantize_fp8_matmul_input(input)
+
+    if groups == 1:
+        result = torch._scaled_mm(input, weight, scale_a=input_scale, scale_b=scale, bias=bias, out_dtype=return_dtype).reshape(mm_output_shape)
+    else:
+        scale = scale.reshape(groups, 1, scale.shape[1] // groups)
+        input_scale = input_scale.reshape(groups, input_scale.shape[0] // groups, 1)
+        weight = weight.reshape(weight.shape[0], groups, weight.shape[1] // groups).transpose(0,1)
+        input = input.reshape(input.shape[0], groups, input.shape[1] // groups).transpose(0,1)
+        result = []
+        if bias is not None:
+            bias = bias.reshape(groups, bias.shape[0] // groups)
+            for i in range(groups):
+                result.append(torch._scaled_mm(input[i], weight[i], scale_a=input_scale[i], scale_b=scale[i], bias=bias[i], out_dtype=return_dtype))
+        else:
+            for i in range(groups):
+                result.append(torch._scaled_mm(input[i], weight[i], scale_a=input_scale[i], scale_b=scale[i], bias=None, out_dtype=return_dtype))
+        result = torch.cat(result, dim=-1).reshape(mm_output_shape)
+
+    if conv_type == 1:
+        result = result.transpose(1,2)
+    elif conv_type == 2:
+        result = result.permute(0,3,1,2)
+    elif conv_type == 3:
+        result = result.permute(0,4,1,2,3)
+    return result
+
+
+def conv_fp8_matmul_tensorwise(
+    input: torch.FloatTensor,
+    weight: torch.Tensor,
+    bias: torch.FloatTensor,
+    scale: torch.FloatTensor,
+    result_shape: torch.Size,
+    reversed_padding_repeated_twice: List[int],
+    padding_mode: str, conv_type: int,
+    groups: int, stride: List[int],
+    padding: List[int], dilation: List[int],
+) -> torch.FloatTensor:
+    return_dtype = input.dtype
+    input, mm_output_shape = process_conv_input(conv_type, input, reversed_padding_repeated_twice, padding_mode, result_shape, stride, padding, dilation)
+    input, scale = quantize_fp8_matmul_input_tensorwise(input, scale)
+    dummy_input_scale = torch.ones(1, device=input.device, dtype=torch.float32)
+
+    if groups == 1:
+        result = torch._scaled_mm(input, weight, scale_a=dummy_input_scale, scale_b=dummy_input_scale, bias=None, out_dtype=scale.dtype)
+    else:
+        weight = weight.reshape(weight.shape[0], groups, weight.shape[1] // groups).transpose(0,1)
+        input = input.reshape(input.shape[0], groups, input.shape[1] // groups).transpose(0,1)
+        result = []
+        for i in range(groups):
+            result.append(torch._scaled_mm(input[i], weight[i], scale_a=dummy_input_scale, scale_b=dummy_input_scale, bias=None, out_dtype=scale.dtype))
+        result = torch.cat(result, dim=-1)
+    if bias is not None:
+        dequantize_symmetric_with_bias(result, bias, scale, return_dtype, mm_output_shape)
+    else:
+        dequantize_symmetric(result, scale, return_dtype, mm_output_shape)
+
+    if conv_type == 1:
+        result = result.transpose(1,2)
+    elif conv_type == 2:
+        result = result.permute(0,3,1,2)
+    elif conv_type == 3:
+        result = result.permute(0,4,1,2,3)
+    return result
+
+
+def conv_int8_matmul(
+    input: torch.FloatTensor,
+    weight: torch.CharTensor,
+    bias: torch.FloatTensor,
+    scale: torch.FloatTensor,
+    result_shape: torch.Size,
+    quantized_weight_shape: torch.Size,
+    weights_dtype: str,
+    reversed_padding_repeated_twice: List[int],
+    padding_mode: str, conv_type: int,
+    groups: int, stride: List[int],
+    padding: List[int], dilation: List[int],
+) -> torch.FloatTensor:
+    return_dtype = input.dtype
+    input, mm_output_shape = process_conv_input(conv_type, input, reversed_padding_repeated_twice, padding_mode, result_shape, stride, padding, dilation)
+    input, scale = quantize_int8_matmul_input(input, scale)
+    if quantized_weight_shape is not None:
+        weight = unpack_int_symetric(weight, quantized_weight_shape, weights_dtype, dtype=torch.int8, transpose=True)
+
+    if groups == 1:
+        result = torch._int_mm(input, weight)
+    else:
+        weight = weight.reshape(weight.shape[0], groups, weight.shape[1] // groups).transpose(0,1)
+        input = input.reshape(input.shape[0], groups, input.shape[1] // groups).transpose(0,1)
+        result = []
+        for i in range(groups):
+            result.append(torch._int_mm(input[i], weight[i]))
+        result = torch.cat(result, dim=-1)
+    if bias is not None:
+        result = dequantize_symmetric_with_bias(result, bias, scale, return_dtype, mm_output_shape)
+    else:
+        result = dequantize_symmetric(result, scale, return_dtype, mm_output_shape)
+
+    if conv_type == 1:
+        result = result.transpose(1,2)
+    elif conv_type == 2:
+        result = result.permute(0,3,1,2)
+    elif conv_type == 3:
+        result = result.permute(0,4,1,2,3)
+    return result
+
+
+def quantized_linear_forward_fp8_matmul(self, input: torch.FloatTensor) -> torch.FloatTensor:
+    if torch.numel(input) / input.shape[-1] < 32:
+        return torch.nn.functional.linear(input, self.sdnq_dequantizer(self.weight, skip_quantized_matmul=True), self.bias)
+    return fp8_matmul(input, self.weight, self.bias, self.sdnq_dequantizer.scale)
+
+
+def quantized_linear_forward_fp8_matmul_tensorwise(self, input: torch.FloatTensor) -> torch.FloatTensor:
+    if torch.numel(input) / input.shape[-1] < 32:
+        return torch.nn.functional.linear(input, self.sdnq_dequantizer(self.weight, skip_quantized_matmul=True), self.bias)
+    return fp8_matmul_tensorwise(input, self.weight, self.bias, self.sdnq_dequantizer.scale)
+
+
+def quantized_linear_forward_int8_matmul(self, input: torch.FloatTensor) -> torch.FloatTensor:
+    if torch.numel(input) / input.shape[-1] < 32:
+        return torch.nn.functional.linear(input, self.sdnq_dequantizer(self.weight, skip_quantized_matmul=True), self.bias)
+    return int8_matmul(input, self.weight, self.bias, self.sdnq_dequantizer.scale, getattr(self.sdnq_dequantizer, "quantized_weight_shape", None), self.sdnq_dequantizer.weights_dtype)
+
+
+def quantized_linear_forward(self, input: torch.FloatTensor) -> torch.FloatTensor:
+    return torch.nn.functional.linear(input, self.sdnq_dequantizer(self.weight), self.bias)
+
+
+def get_conv_args(input_ndim: int, stride, padding, dilation):
+    if input_ndim == 3:
+        conv_type = 1
+    elif input_ndim == 4:
+        conv_type = 2
+    else:
+        conv_type = 3
+    if isinstance(stride, int):
+        stride = (stride,) * conv_type
+    if isinstance(padding, int):
+        padding = (padding,) * conv_type
+    if isinstance(dilation, int):
+        dilation = (dilation,) * conv_type
+    if conv_type == 1:
+        stride = (1, stride[0])
+        padding = (0, padding[0])
+        dilation = (1, dilation[0])
+    return conv_type, stride, padding, dilation
+
+
+def quantized_conv_forward_fp8_matmul(self, input) -> torch.FloatTensor:
+    if torch.numel(input) / input.shape[2] < 32:
+        return self._conv_forward(input, self.sdnq_dequantizer(self.weight, skip_quantized_matmul=True), self.bias)
+    conv_type, stride, padding, dilation = get_conv_args(input.ndim, self.stride, self.padding, self.dilation)
+    return conv_fp8_matmul(
+        input, self.weight, self.bias,
+        self.sdnq_dequantizer.scale,
+        self.sdnq_dequantizer.result_shape,
+        self._reversed_padding_repeated_twice,
+        self.padding_mode, conv_type,
+        self.groups, stride, padding, dilation,
+    )
+
+
+def quantized_conv_forward_fp8_matmul_tensorwise(self, input) -> torch.FloatTensor:
+    if torch.numel(input) / input.shape[2] < 32:
+        return self._conv_forward(input, self.sdnq_dequantizer(self.weight, skip_quantized_matmul=True), self.bias)
+    conv_type, stride, padding, dilation = get_conv_args(input.ndim, self.stride, self.padding, self.dilation)
+    return conv_fp8_matmul_tensorwise(
+        input, self.weight, self.bias,
+        self.sdnq_dequantizer.scale,
+        self.sdnq_dequantizer.result_shape,
+        self._reversed_padding_repeated_twice,
+        self.padding_mode, conv_type,
+        self.groups, stride, padding, dilation,
+    )
+
+
+def quantized_conv_forward_int8_matmul(self, input) -> torch.FloatTensor:
+    if torch.numel(input) / input.shape[2] < 32:
+        return self._conv_forward(input, self.sdnq_dequantizer(self.weight, skip_quantized_matmul=True), self.bias)
+    conv_type, stride, padding, dilation = get_conv_args(input.ndim, self.stride, self.padding, self.dilation)
+    return conv_int8_matmul(
+        input, self.weight, self.bias,
+        self.sdnq_dequantizer.scale,
+        self.sdnq_dequantizer.result_shape,
+        getattr(self.sdnq_dequantizer, "quantized_weight_shape", None),
+        self.sdnq_dequantizer.weights_dtype,
+        self._reversed_padding_repeated_twice,
+        self.padding_mode, conv_type,
+        self.groups, stride, padding, dilation,
+    )
+
+
+def quantized_conv_forward(self, input) -> torch.FloatTensor:
+    return self._conv_forward(input, self.sdnq_dequantizer(self.weight), self.bias)
+
+
+def quantized_conv_transpose_1d_forward(self, input: torch.FloatTensor, output_size: Optional[list[int]] = None) -> torch.FloatTensor:
+    output_padding = self._output_padding(input, output_size, self.stride, self.padding, self.kernel_size, 1, self.dilation)
+    return torch.nn.functional.conv_transpose1d(input, self.sdnq_dequantizer(self.weight), self.bias, self.stride, self.padding, output_padding, self.groups, self.dilation)
+
+
+def quantized_conv_transpose_2d_forward(self, input: torch.FloatTensor, output_size: Optional[list[int]] = None) -> torch.FloatTensor:
+    output_padding = self._output_padding(input, output_size, self.stride, self.padding, self.kernel_size, 2, self.dilation)
+    return torch.nn.functional.conv_transpose2d(input, self.sdnq_dequantizer(self.weight), self.bias, self.stride, self.padding, output_padding, self.groups, self.dilation)
+
+
+def quantized_conv_transpose_3d_forward(self, input: torch.FloatTensor, output_size: Optional[list[int]] = None) -> torch.FloatTensor:
+    output_padding = self._output_padding(input, output_size, self.stride, self.padding, self.kernel_size, 3, self.dilation)
+    return torch.nn.functional.conv_transpose3d(input, self.sdnq_dequantizer(self.weight), self.bias, self.stride, self.padding, output_padding, self.groups, self.dilation)
+
+
+if shared.opts.sdnq_dequantize_compile:
+    try:
+        torch._dynamo.config.cache_size_limit = max(8192, torch._dynamo.config.cache_size_limit)
+        int8_matmul = torch.compile(int8_matmul, fullgraph=True)
+        fp8_matmul = torch.compile(fp8_matmul, fullgraph=True)
+        fp8_matmul_tensorwise = torch.compile(fp8_matmul_tensorwise, fullgraph=True)
+        conv_int8_matmul = torch.compile(conv_int8_matmul, fullgraph=True)
+        conv_fp8_matmul = torch.compile(conv_fp8_matmul, fullgraph=True)
+        conv_fp8_matmul_tensorwise = torch.compile(conv_fp8_matmul_tensorwise, fullgraph=True)
+    except Exception as e:
+        shared.log.warning(f"Quantization: type=sdnq MatMul using torch.compile is not available: {e}")

modules/sdnq/packed_int.py

@@ -0,0 +1,268 @@
+# pylint: disable=redefined-builtin,no-member,protected-access
+
+from typing import Optional
+import torch
+
+from .common import dtype_dict
+
+
+def pack_int_symetric(tensor: torch.ByteTensor, weights_dtype: str) -> torch.ByteTensor:
+    return packed_int_function_dict[weights_dtype]["pack"](tensor.to(dtype=dtype_dict[weights_dtype]["torch_dtype"]).sub_(dtype_dict[weights_dtype]["min"]).to(dtype=dtype_dict[weights_dtype]["storage_dtype"]))
+
+
+def unpack_int_symetric(packed_tensor: torch.CharTensor, shape: torch.Size, weights_dtype: str, dtype: Optional[torch.dtype] = None, transpose: Optional[bool] = False) -> torch.ByteTensor:
+    if dtype is None:
+        dtype = dtype_dict[weights_dtype]["torch_dtype"]
+    result = packed_int_function_dict[weights_dtype]["unpack"](packed_tensor, shape).to(dtype=dtype).add_(dtype_dict[weights_dtype]["min"])
+    if transpose:
+        result = result.transpose(0,1)
+    return result
+
+
+def pack_uint7(tensor: torch.Tensor) -> torch.Tensor:
+    if tensor.dtype != torch.uint8:
+        raise RuntimeError(f"Invalid tensor dtype {tensor.type}. torch.uint8 type is supported.")
+    packed_tensor = tensor.contiguous().reshape(-1, 8)
+    packed_tensor = torch.stack(
+        (
+            torch.bitwise_or(packed_tensor[:, 0], torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 7], 1), 128)),
+            torch.bitwise_or(packed_tensor[:, 1], torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 7], 2), 128)),
+            torch.bitwise_or(packed_tensor[:, 2], torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 7], 3), 128)),
+            torch.bitwise_or(packed_tensor[:, 3], torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 7], 4), 128)),
+            torch.bitwise_or(packed_tensor[:, 4], torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 7], 5), 128)),
+            torch.bitwise_or(packed_tensor[:, 5], torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 7], 6), 128)),
+            torch.bitwise_or(packed_tensor[:, 6], torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 7], 7), 128)),
+        ),
+        dim=-1
+    )
+    return packed_tensor
+
+
+def pack_uint6(tensor: torch.Tensor) -> torch.Tensor:
+    if tensor.dtype != torch.uint8:
+        raise RuntimeError(f"Invalid tensor dtype {tensor.type}. torch.uint8 type is supported.")
+    packed_tensor = tensor.contiguous().reshape(-1, 4)
+    packed_tensor = torch.stack(
+        (
+            torch.bitwise_or(packed_tensor[:, 0], torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 3], 2), 192)),
+            torch.bitwise_or(packed_tensor[:, 1], torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 3], 4), 192)),
+            torch.bitwise_or(packed_tensor[:, 2], torch.bitwise_left_shift(packed_tensor[:, 3], 6)),
+        ),
+        dim=-1
+    )
+    return packed_tensor
+
+
+def pack_uint5(tensor: torch.Tensor) -> torch.Tensor:
+    if tensor.dtype != torch.uint8:
+        raise RuntimeError(f"Invalid tensor dtype {tensor.type}. torch.uint8 type is supported.")
+    packed_tensor = tensor.contiguous().reshape(-1, 8)
+    packed_tensor = torch.stack(
+        (
+            torch.bitwise_or(packed_tensor[:, 0], torch.bitwise_left_shift(packed_tensor[:, 5], 5)),
+            torch.bitwise_or(packed_tensor[:, 1], torch.bitwise_left_shift(packed_tensor[:, 6], 5)),
+            torch.bitwise_or(packed_tensor[:, 2], torch.bitwise_left_shift(packed_tensor[:, 7], 5)),
+            torch.bitwise_or(
+                packed_tensor[:, 3],
+                torch.bitwise_or(
+                    torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 5], 2), 96),
+                    torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 7], 3), 128),
+                ),
+            ),
+            torch.bitwise_or(
+                packed_tensor[:, 4],
+                torch.bitwise_or(
+                    torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 6], 2), 96),
+                    torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 7], 4), 128),
+                ),
+            ),
+        ),
+        dim=-1
+    )
+    return packed_tensor
+
+
+def pack_uint4(tensor: torch.Tensor) -> torch.Tensor:
+    if tensor.dtype != torch.uint8:
+        raise RuntimeError(f"Invalid tensor dtype {tensor.type}. torch.uint8 type is supported.")
+    packed_tensor = tensor.contiguous().reshape(-1, 2)
+    packed_tensor = torch.bitwise_or(packed_tensor[:, 0], torch.bitwise_left_shift(packed_tensor[:, 1], 4))
+    return packed_tensor
+
+
+def pack_uint3(tensor: torch.Tensor) -> torch.Tensor:
+    if tensor.dtype != torch.uint8:
+        raise RuntimeError(f"Invalid tensor dtype {tensor.type}. torch.uint8 type is supported.")
+    packed_tensor = tensor.contiguous().reshape(-1, 8)
+    packed_tensor = torch.stack(
+        (
+            torch.bitwise_or(
+                torch.bitwise_or(packed_tensor[:, 0], torch.bitwise_left_shift(packed_tensor[:, 1], 3)),
+                torch.bitwise_left_shift(packed_tensor[:, 6], 6),
+            ),
+            torch.bitwise_or(
+                torch.bitwise_or(packed_tensor[:, 2], torch.bitwise_left_shift(packed_tensor[:, 3], 3)),
+                torch.bitwise_left_shift(packed_tensor[:, 7], 6),
+            ),
+            torch.bitwise_or(
+                torch.bitwise_or(packed_tensor[:, 4], torch.bitwise_left_shift(packed_tensor[:, 5], 3)),
+                torch.bitwise_or(
+                    torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 6], 4), 64),
+                    torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 7], 5), 128),
+                )
+            ),
+        ),
+        dim=-1
+    )
+    return packed_tensor
+
+
+def pack_uint2(tensor: torch.Tensor) -> torch.Tensor:
+    if tensor.dtype != torch.uint8:
+        raise RuntimeError(f"Invalid tensor dtype {tensor.type}. torch.uint8 type is supported.")
+    packed_tensor = tensor.contiguous().reshape(-1, 4)
+    packed_tensor = torch.bitwise_or(
+        torch.bitwise_or(packed_tensor[:, 0], torch.bitwise_left_shift(packed_tensor[:, 1], 2)),
+        torch.bitwise_or(torch.bitwise_left_shift(packed_tensor[:, 2], 4), torch.bitwise_left_shift(packed_tensor[:, 3], 6)),
+    )
+    return packed_tensor
+
+
+def unpack_uint7(packed_tensor: torch.Tensor, shape: torch.Size) -> torch.Tensor:
+    result = torch.stack(
+        (
+            torch.bitwise_and(packed_tensor[:, 0], 127),
+            torch.bitwise_and(packed_tensor[:, 1], 127),
+            torch.bitwise_and(packed_tensor[:, 2], 127),
+            torch.bitwise_and(packed_tensor[:, 3], 127),
+            torch.bitwise_and(packed_tensor[:, 4], 127),
+            torch.bitwise_and(packed_tensor[:, 5], 127),
+            torch.bitwise_and(packed_tensor[:, 6], 127),
+            torch.bitwise_or(
+                torch.bitwise_or(
+                    torch.bitwise_or(
+                        torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 0], 1), 64),
+                        torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 1], 2), 32),
+                    ),
+                    torch.bitwise_or(
+                        torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 2], 3), 16),
+                        torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 3], 4), 8),
+                    ),
+                ),
+                torch.bitwise_or(
+                    torch.bitwise_or(
+                        torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 4], 5), 4),
+                        torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 5], 6), 2),
+                    ),
+                    torch.bitwise_right_shift(packed_tensor[:, 6], 7),
+                ),
+            )
+        ),
+        dim=-1
+    ).reshape(shape)
+    return result
+
+
+def unpack_uint6(packed_tensor: torch.Tensor, shape: torch.Size) -> torch.Tensor:
+    result = torch.stack(
+        (
+            torch.bitwise_and(packed_tensor[:, 0], 63),
+            torch.bitwise_and(packed_tensor[:, 1], 63),
+            torch.bitwise_and(packed_tensor[:, 2], 63),
+            torch.bitwise_or(
+                torch.bitwise_or(
+                    torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 0], 2), 48),
+                    torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 1], 4), 12),
+                ),
+                torch.bitwise_right_shift(packed_tensor[:, 2], 6)
+            )
+        ),
+        dim=-1
+    ).reshape(shape)
+    return result
+
+
+def unpack_uint5(packed_tensor: torch.Tensor, shape: torch.Size) -> torch.Tensor:
+    result = torch.stack(
+        (
+            torch.bitwise_and(packed_tensor[:, 0], 31),
+            torch.bitwise_and(packed_tensor[:, 1], 31),
+            torch.bitwise_and(packed_tensor[:, 2], 31),
+            torch.bitwise_and(packed_tensor[:, 3], 31),
+            torch.bitwise_and(packed_tensor[:, 4], 31),
+            torch.bitwise_or(
+                torch.bitwise_right_shift(packed_tensor[:, 0], 5),
+                torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 3], 2), 24),
+            ),
+            torch.bitwise_or(
+                torch.bitwise_right_shift(packed_tensor[:, 1], 5),
+                torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 4], 2), 24),
+            ),
+            torch.bitwise_or(
+                torch.bitwise_right_shift(packed_tensor[:, 2], 5),
+                torch.bitwise_or(
+                    torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 3], 3), 16),
+                    torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 4], 4), 8),
+                ),
+            ),
+        ),
+        dim=-1
+    ).reshape(shape)
+    return result
+
+
+def unpack_uint4(packed_tensor: torch.Tensor, shape: torch.Size) -> torch.Tensor:
+    result = torch.stack((torch.bitwise_and(packed_tensor, 15), torch.bitwise_right_shift(packed_tensor, 4)), dim=-1).reshape(shape)
+    return result
+
+
+def unpack_uint3(packed_tensor: torch.Tensor, shape: torch.Size) -> torch.Tensor:
+    result = torch.stack(
+        (
+            torch.bitwise_and(packed_tensor[:, 0], 7),
+            torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 0], 3), 7),
+            torch.bitwise_and(packed_tensor[:, 1], 7),
+            torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 1], 3), 7),
+            torch.bitwise_and(packed_tensor[:, 2], 7),
+            torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 2], 3), 7),
+            torch.bitwise_or(
+                torch.bitwise_right_shift(packed_tensor[:, 0], 6),
+                torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 2], 4), 4),
+            ),
+            torch.bitwise_or(
+                torch.bitwise_right_shift(packed_tensor[:, 1], 6),
+                torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 2], 5), 4),
+            ),
+        ),
+        dim=-1
+    ).reshape(shape)
+    return result
+
+
+def unpack_uint2(packed_tensor: torch.Tensor, shape: torch.Size) -> torch.Tensor:
+    result = torch.stack(
+        (
+            torch.bitwise_and(packed_tensor, 3),
+            torch.bitwise_and(torch.bitwise_right_shift(packed_tensor, 2), 3),
+            torch.bitwise_and(torch.bitwise_right_shift(packed_tensor, 4), 3),
+            torch.bitwise_right_shift(packed_tensor, 6),
+        ),
+        dim=-1
+    ).reshape(shape)
+    return result
+
+
+packed_int_function_dict = {
+    "int7": {"pack": pack_uint7, "unpack": unpack_uint7},
+    "int6": {"pack": pack_uint6, "unpack": unpack_uint6},
+    "int5": {"pack": pack_uint5, "unpack": unpack_uint5},
+    "int4": {"pack": pack_uint4, "unpack": unpack_uint4},
+    "int3": {"pack": pack_uint3, "unpack": unpack_uint3},
+    "int2": {"pack": pack_uint2, "unpack": unpack_uint2},
+    "uint7": {"pack": pack_uint7, "unpack": unpack_uint7},
+    "uint6": {"pack": pack_uint6, "unpack": unpack_uint6},
+    "uint5": {"pack": pack_uint5, "unpack": unpack_uint5},
+    "uint4": {"pack": pack_uint4, "unpack": unpack_uint4},
+    "uint3": {"pack": pack_uint3, "unpack": unpack_uint3},
+    "uint2": {"pack": pack_uint2, "unpack": unpack_uint2},
+}

modules/shared.py

@@ -353,7 +353,7 @@ def get_default_modes():
                 default_offload_mode = "sequential"
                 default_diffusers_offload_min_gpu_memory = 0
                 log.info(f"Device detect: memory={gpu_memory:.1f} default=sequential optimization=lowvram")
-            elif gpu_memory <= 8:
+            elif gpu_memory <= 12:
                 cmd_opts.medvram = True # VAE Tiling and other stuff
                 default_offload_mode = "balanced"
                 default_diffusers_offload_min_gpu_memory = 0
@@ -518,13 +518,15 @@ options_templates.update(options_section(("quantization", "Quantization Settings
     "sdnq_quantize_sep": OptionInfo("<h2>SDNQ: SD.Next Quantization</h2>", "", gr.HTML),
     "sdnq_quantize_weights": OptionInfo([], "Quantization enabled", gr.CheckboxGroup, {"choices": ["Model", "Transformer", "VAE", "TE", "Video", "LLM", "ControlNet"], "visible": native}),
     "sdnq_quantize_mode": OptionInfo("pre", "Quantization mode", gr.Dropdown, {"choices": ["pre", "post"], "visible": native}),
-    "sdnq_quantize_weights_mode": OptionInfo("int8", "Quantization type", gr.Dropdown, {"choices": ["int8", "int6", "uint4", "float8_e4m3fn", "uint8", "uint6", "int4", "float8_e5m2", "float8_e4m3fnuz", "float8_e5m2fnuz", "int2", "uint2", "uint1"], "visible": native}),
+    "sdnq_quantize_weights_mode": OptionInfo("int8", "Quantization type", gr.Dropdown, {"choices": ["int8", "float8_e4m3fn", "int7", "int6", "int5", "uint4", "uint3", "uint2", "float8_e5m2", "float8_e4m3fnuz", "float8_e5m2fnuz", "uint8", "uint7", "uint6", "uint5", "int4", "int3", "int2", "uint1"], "visible": native}),
+    "sdnq_quantize_weights_mode_te": OptionInfo("default", "Quantization type for Text Encoders", gr.Dropdown, {"choices": ["default", "int8", "float8_e4m3fn", "int7", "int6", "int5", "uint4", "uint3", "uint2", "float8_e5m2", "float8_e4m3fnuz", "float8_e5m2fnuz", "uint8", "uint7", "uint6", "uint5", "int4", "int3", "int2", "uint1"], "visible": native}),
     "sdnq_quantize_weights_group_size": OptionInfo(0, "Group size", gr.Slider, {"minimum": -1, "maximum": 4096, "step": 1, "visible": native}),
     "sdnq_quantize_conv_layers": OptionInfo(False, "Quantize the convolutional layers", gr.Checkbox, {"visible": native}),
-    "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_dequantize_compile": OptionInfo(devices.has_triton(), "Dequantize using torch.compile", gr.Checkbox, {"visible": native}),
     "sdnq_use_quantized_matmul": OptionInfo(False, "Use Quantized MatMul", gr.Checkbox, {"visible": native}),
+    "sdnq_use_quantized_matmul_conv": OptionInfo(False, "Use Quantized MatMul with convolutional layers", gr.Checkbox, {"visible": native}),
     "sdnq_quantize_with_gpu": OptionInfo(True, "Quantize with the GPU", gr.Checkbox, {"visible": native}),
+    "sdnq_dequantize_fp32": OptionInfo(False, "Dequantize using full precision", 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),

modules/shared_items.py

@@ -29,7 +29,7 @@ pipelines = {
     'FLEX': getattr(diffusers, 'AutoPipelineForText2Image', None),
     'Sana': getattr(diffusers, 'SanaPipeline', None),
     'Lumina-Next': getattr(diffusers, 'LuminaText2ImgPipeline', None),
-    'Lumina 2': getattr(diffusers, 'Lumina2Text2ImgPipeline', None),
+    'Lumina 2': getattr(diffusers, 'Lumina2Pipeline', None),
     'AuraFlow': getattr(diffusers, 'AuraFlowPipeline', None),
     'Kandinsky 2.1': getattr(diffusers, 'KandinskyCombinedPipeline', None),
     'Kandinsky 2.2': getattr(diffusers, 'KandinskyV22CombinedPipeline', None),

repositories/codeformer/basicsr/__init__.py

@@ -8,4 +8,5 @@ from .models import *
 from .ops import *
 from .train import *
 from .utils import *
-from .version import __gitsha__, __version__
+__gitsha__ = '366a46c91d51923c56e09963dbc358bc61315408'
+__version__ = '1.3.2'

repositories/codeformer/basicsr/losses/losses.py

@@ -1,5 +1,4 @@
 import math
-import lpips
 import torch
 from torch import autograd as autograd
 from torch import nn as nn
@@ -260,6 +259,7 @@ class LPIPSLoss(nn.Module):
             use_input_norm=True,
             range_norm=False,):
         super(LPIPSLoss, self).__init__()
+        import lpips
         self.perceptual = lpips.LPIPS(net="vgg", spatial=False).eval()
         self.loss_weight = loss_weight
         self.use_input_norm = use_input_norm

requirements.txt

@@ -33,7 +33,7 @@ pi-heif
 # versioned
 rich==14.0.0
 safetensors==0.5.3
-tensordict==0.1.2
+tensordict==0.8.3
 peft==0.15.2
 httpx==0.24.1
 compel==2.0.3
@@ -45,14 +45,15 @@ accelerate==1.7.0
 opencv-contrib-python-headless==4.9.0.80
 einops==0.4.1
 gradio==3.43.2
-huggingface_hub==0.31.2
+huggingface_hub==0.33.0
 numexpr==2.10.2
 numpy==1.26.4
+pandas==2.3.0
 numba==0.61.2
 protobuf==4.25.3
 pytorch_lightning==1.9.4
 tokenizers==0.21.1
-transformers==4.52.3
+transformers==4.52.4
 urllib3==1.26.19
 Pillow==10.4.0
 timm==0.9.16
@@ -63,7 +64,6 @@ typing-extensions==4.12.2
 # additional
 blendmodes
 scipy
-pandas
 torchdiffeq
 dctorch
 scikit-image