diff --git a/CHANGELOG.md b/CHANGELOG.md
index 968bb8e1c..c2f459f35 100644
--- a/CHANGELOG.md
+++ b/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
diff --git a/installer.py b/installer.py
index cecf6d0c6..bf4d0b408 100644
--- a/installer.py
+++ b/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
diff --git a/modules/devices.py b/modules/devices.py
index 8bf2def3f..1c35f2683 100644
--- a/modules/devices.py
+++ b/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:
diff --git a/modules/intel/ipex/__init__.py b/modules/intel/ipex/__init__.py
index 369367ef8..a44531f35 100644
--- a/modules/intel/ipex/__init__.py
+++ b/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
diff --git a/modules/intel/ipex/diffusers.py b/modules/intel/ipex/diffusers.py
index 033b74cbe..d3487fefd 100644
--- a/modules/intel/ipex/diffusers.py
+++ b/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
diff --git a/modules/intel/ipex/hijacks.py b/modules/intel/ipex/hijacks.py
index 0ce8abdc5..d81d7b05c 100644
--- a/modules/intel/ipex/hijacks.py
+++ b/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
diff --git a/modules/lora/lora_apply.py b/modules/lora/lora_apply.py
index 7edc9f461..b971541cd 100644
--- a/modules/lora/lora_apply.py
+++ b/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
diff --git a/modules/meissonic/scheduler.py b/modules/meissonic/scheduler.py
index 3d2fe4276..757469b4d 100644
--- a/modules/meissonic/scheduler.py
+++ b/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
)
diff --git a/modules/model_flux.py b/modules/model_flux.py
index 17ebe06c7..dab00d86a 100644
--- a/modules/model_flux.py
+++ b/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:
diff --git a/modules/model_lumina.py b/modules/model_lumina.py
index f9d3b9abd..5f46da24f 100644
--- a/modules/model_lumina.py
+++ b/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,
diff --git a/modules/model_pixart.py b/modules/model_pixart.py
index 58d3dbfa7..6f6d6cf1c 100644
--- a/modules/model_pixart.py
+++ b/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,
diff --git a/modules/model_quant.py b/modules/model_quant.py
index 96579c4eb..3c57ac67c 100644
--- a/modules/model_quant.py
+++ b/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'
diff --git a/modules/model_quant_sdnq.py b/modules/model_quant_sdnq.py
deleted file mode 100644
index 1281b05e4..000000000
--- a/modules/model_quant_sdnq.py
+++ /dev/null
@@ -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
diff --git a/modules/modeldata.py b/modules/modeldata.py
index b7be0866b..688cff23d 100644
--- a/modules/modeldata.py
+++ b/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
diff --git a/modules/omnigen/transformer.py b/modules/omnigen/transformer.py
index f3bcdb15a..d166309ca 100644
--- a/modules/omnigen/transformer.py
+++ b/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]
diff --git a/modules/postprocess/gfpgan_model.py b/modules/postprocess/gfpgan_model.py
index ad0aa8221..8eb998921 100644
--- a/modules/postprocess/gfpgan_model.py
+++ b/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
diff --git a/modules/processing_vae.py b/modules/processing_vae.py
index 5a375adc2..36d5d0dda 100644
--- a/modules/processing_vae.py
+++ b/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]
diff --git a/modules/rocm.py b/modules/rocm.py
index f16809291..816a3b10c 100644
--- a/modules/rocm.py
+++ b/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
diff --git a/modules/sd_models.py b/modules/sd_models.py
index 22ba21a49..0bea615ff 100644
--- a/modules/sd_models.py
+++ b/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
diff --git a/modules/sd_models_utils.py b/modules/sd_models_utils.py
index 0e1b72f4c..546297d25 100644
--- a/modules/sd_models_utils.py
+++ b/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'):
diff --git a/modules/sd_vae_taesd.py b/modules/sd_vae_taesd.py
index da12422f1..3bcd1b322 100644
--- a/modules/sd_vae_taesd.py
+++ b/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")
diff --git a/modules/sdnq/__init__.py b/modules/sdnq/__init__.py
new file mode 100644
index 000000000..f5b7179a7
--- /dev/null
+++ b/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
diff --git a/modules/sdnq/common.py b/modules/sdnq/common.py
new file mode 100644
index 000000000..ccfdf1227
--- /dev/null
+++ b/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
diff --git a/modules/sdnq/dequantizer.py b/modules/sdnq/dequantizer.py
new file mode 100644
index 000000000..be59019a3
--- /dev/null
+++ b/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
diff --git a/modules/sdnq/forward.py b/modules/sdnq/forward.py
new file mode 100644
index 000000000..9caa12f7d
--- /dev/null
+++ b/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}")
diff --git a/modules/sdnq/packed_int.py b/modules/sdnq/packed_int.py
new file mode 100644
index 000000000..e4717086e
--- /dev/null
+++ b/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},
+}
diff --git a/modules/shared.py b/modules/shared.py
index a1c089c48..2231e0484 100644
--- a/modules/shared.py
+++ b/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("SDNQ: SD.Next Quantization
", "", 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("BitsAndBytes
", "", gr.HTML),
diff --git a/modules/shared_items.py b/modules/shared_items.py
index d8796daf5..9a09b8c7f 100644
--- a/modules/shared_items.py
+++ b/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),
diff --git a/repositories/codeformer/basicsr/__init__.py b/repositories/codeformer/basicsr/__init__.py
index c7ffcccd7..24be6f0af 100644
--- a/repositories/codeformer/basicsr/__init__.py
+++ b/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'
diff --git a/repositories/codeformer/basicsr/losses/losses.py b/repositories/codeformer/basicsr/losses/losses.py
index 1bcf272cf..71331aa01 100644
--- a/repositories/codeformer/basicsr/losses/losses.py
+++ b/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
diff --git a/requirements.txt b/requirements.txt
index 5c464a7a1..f5eb92014 100644
--- a/requirements.txt
+++ b/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
diff --git a/wiki b/wiki
index 693b0dafa..34e99de10 160000
--- a/wiki
+++ b/wiki
@@ -1 +1 @@
-Subproject commit 693b0dafa0dea17fe8057201f397e80a9b46b0c0
+Subproject commit 34e99de10210375593daec696f1b651d82ef0cf0