From 251945939b0ca413105e9db63addd6c1d3a80cef Mon Sep 17 00:00:00 2001 From: Dity0 Date: Mon, 10 Aug 2026 22:51:22 +0300 Subject: [PATCH] Add Lloyd-Max quantization with use_codebook --- modules/lora/lora_apply.py | 2 ++ modules/model_quant.py | 9 +++++++-- modules/ui_definitions.py | 2 ++ 3 files changed, 11 insertions(+), 2 deletions(-) diff --git a/modules/lora/lora_apply.py b/modules/lora/lora_apply.py index 467330e3e..f8353320e 100644 --- a/modules/lora/lora_apply.py +++ b/modules/lora/lora_apply.py @@ -210,8 +210,10 @@ def network_add_weights(self: torch.nn.Conv2d | torch.nn.Linear | torch.nn.Group use_quantized_matmul_conv=sdnq_dequantizer.use_quantized_matmul, use_svd=use_svd, use_hadamard=sdnq_dequantizer.use_hadamard, + use_codebook=sdnq_dequantizer.use_codebook, dequantize_fp32=dequantize_fp32, svd_steps=shared.opts.sdnq_svd_steps, + codebook_steps=shared.opts.sdnq_codebook_steps, quant_conv=True, # quant_conv is True if conv layers ends up here non_blocking=False, quantization_device=devices.device, diff --git a/modules/model_quant.py b/modules/model_quant.py index 031e986cb..a748acb37 100644 --- a/modules/model_quant.py +++ b/modules/model_quant.py @@ -159,9 +159,11 @@ def create_sdnq_config(kwargs = None, group_size=shared.opts.sdnq_group_size, svd_rank=shared.opts.sdnq_svd_rank, svd_steps=shared.opts.sdnq_svd_steps, + codebook_steps=shared.opts.sdnq_codebook_steps, dynamic_loss_threshold=shared.opts.sdnq_dynamic_loss_threshold, use_svd=shared.opts.sdnq_use_svd, use_hadamard=shared.opts.sdnq_use_hadamard, + use_codebook=shared.opts.sdnq_use_codebook, quant_conv=shared.opts.sdnq_quantize_conv_layers, quant_embedding=shared.opts.sdnq_quantize_embedding_layers, use_quantized_matmul=use_quantized_matmul, @@ -176,7 +178,8 @@ def create_sdnq_config(kwargs = None, ) svd = f'{shared.opts.sdnq_use_svd} rank={shared.opts.sdnq_svd_rank} steps={shared.opts.sdnq_svd_steps}' if shared.opts.sdnq_use_svd else f'{shared.opts.sdnq_use_svd}' hadamard = f'{shared.opts.sdnq_use_hadamard} group={shared.opts.sdnq_hadamard_group_size}' if shared.opts.sdnq_use_hadamard else f'{shared.opts.sdnq_use_hadamard}' - log.debug(f'Quantization: module="{module}" type=sdnq mode=pre dtype={weights_dtype} svd={svd} hadamard={hadamard} dynamic={shared.opts.sdnq_use_dynamic_quantization} group={shared.opts.sdnq_group_size} loss={shared.opts.sdnq_dynamic_loss_threshold} matmul_dtype={quantized_matmul_dtype_log} matmul_quant={use_quantized_matmul} matmul_conv={shared.opts.sdnq_use_quantized_matmul_conv} quant_conv={shared.opts.sdnq_quantize_conv_layers} quant_embed={shared.opts.sdnq_quantize_embedding_layers} fp32={shared.opts.sdnq_dequantize_fp32} device={quantization_device} return={return_device} gpu={shared.opts.sdnq_quantize_with_gpu} map={shared.opts.device_map}') + codebook = f'{shared.opts.sdnq_use_codebook} steps={shared.opts.sdnq_codebook_steps}' if shared.opts.sdnq_use_codebook else f'{shared.opts.sdnq_use_codebook}' + log.debug(f'Quantization: module="{module}" type=sdnq mode=pre dtype={weights_dtype} svd={svd} hadamard={hadamard} codebook={codebook} dynamic={shared.opts.sdnq_use_dynamic_quantization} group={shared.opts.sdnq_group_size} loss={shared.opts.sdnq_dynamic_loss_threshold} matmul_dtype={quantized_matmul_dtype_log} matmul_quant={use_quantized_matmul} matmul_conv={shared.opts.sdnq_use_quantized_matmul_conv} quant_conv={shared.opts.sdnq_quantize_conv_layers} quant_embed={shared.opts.sdnq_quantize_embedding_layers} fp32={shared.opts.sdnq_dequantize_fp32} device={quantization_device} return={return_device} gpu={shared.opts.sdnq_quantize_with_gpu} map={shared.opts.device_map}') if len(modules_to_not_convert) > 0 or modules_dtype_dict: log.debug(f'Quantization: module={module} type=sdnq skip_modules={modules_to_not_convert} modules_dtype_dict={modules_dtype_dict}') if kwargs is None: @@ -430,9 +433,11 @@ def sdnq_quantize_model(model, op=None, sd_model=None, do_gc: bool = True, weigh group_size=shared.opts.sdnq_group_size, svd_rank=shared.opts.sdnq_svd_rank, svd_steps=shared.opts.sdnq_svd_steps, + codebook_steps=shared.opts.sdnq_codebook_steps, dynamic_loss_threshold=shared.opts.sdnq_dynamic_loss_threshold, use_svd=shared.opts.sdnq_use_svd, use_hadamard=shared.opts.sdnq_use_hadamard, + use_codebook=shared.opts.sdnq_use_codebook, quant_conv=shared.opts.sdnq_quantize_conv_layers, quant_embedding=shared.opts.sdnq_quantize_embedding_layers, use_quantized_matmul=use_quantized_matmul, @@ -476,7 +481,7 @@ def sdnq_quantize_model(model, op=None, sd_model=None, do_gc: bool = True, weigh model = model.to(devices.cpu) if do_gc: devices.torch_gc(force=True, reason='sdnq') - log.debug(f'Quantization: module="{op if op is not None else model.__class__}" type=sdnq mode=post dtype={weights_dtype} matmul_dtype={quantized_matmul_dtype_log} matmul={use_quantized_matmul} svd={shared.opts.sdnq_use_svd} hadamard={shared.opts.sdnq_use_hadamard} dynamic={shared.opts.sdnq_use_dynamic_quantization}:group={shared.opts.sdnq_group_size}:hadamard_group={shared.opts.sdnq_hadamard_group_size}:rank={shared.opts.sdnq_svd_rank}:steps={shared.opts.sdnq_svd_steps}:loss={shared.opts.sdnq_dynamic_loss_threshold} matmul_conv={shared.opts.sdnq_use_quantized_matmul_conv} quant_conv={shared.opts.sdnq_quantize_conv_layers} quant_embedding={shared.opts.sdnq_quantize_embedding_layers} fp32={shared.opts.sdnq_dequantize_fp32} gpu={shared.opts.sdnq_quantize_with_gpu} device={quantization_device} return={return_device} map={shared.opts.device_map} non_blocking={shared.opts.diffusers_offload_nonblocking} modules_skip={modules_to_not_convert} modules_dtype={modules_dtype_dict}') + log.debug(f'Quantization: module="{op if op is not None else model.__class__}" type=sdnq mode=post dtype={weights_dtype} matmul_dtype={quantized_matmul_dtype_log} matmul={use_quantized_matmul} svd={shared.opts.sdnq_use_svd} hadamard={shared.opts.sdnq_use_hadamard} codebook={shared.opts.sdnq_use_codebook} dynamic={shared.opts.sdnq_use_dynamic_quantization}:group={shared.opts.sdnq_group_size}:hadamard_group={shared.opts.sdnq_hadamard_group_size}:rank={shared.opts.sdnq_svd_rank}:steps={shared.opts.sdnq_svd_steps}:loss={shared.opts.sdnq_dynamic_loss_threshold} matmul_conv={shared.opts.sdnq_use_quantized_matmul_conv} quant_conv={shared.opts.sdnq_quantize_conv_layers} quant_embedding={shared.opts.sdnq_quantize_embedding_layers} fp32={shared.opts.sdnq_dequantize_fp32} gpu={shared.opts.sdnq_quantize_with_gpu} device={quantization_device} return={return_device} map={shared.opts.device_map} non_blocking={shared.opts.diffusers_offload_nonblocking} modules_skip={modules_to_not_convert} modules_dtype={modules_dtype_dict}') return model diff --git a/modules/ui_definitions.py b/modules/ui_definitions.py index c01e52fb7..5944dc426 100644 --- a/modules/ui_definitions.py +++ b/modules/ui_definitions.py @@ -174,9 +174,11 @@ def create_settings(cmd_opts): "sdnq_hadamard_group_size": OptionInfo(256, "Hadamard group size", gr.Slider, {"minimum": 4, "maximum": 4096, "step": 1}), "sdnq_svd_rank": OptionInfo(32, "SVD rank size", gr.Slider, {"minimum": 1, "maximum": 512, "step": 1}), "sdnq_svd_steps": OptionInfo(8, "SVD steps", gr.Slider, {"minimum": 1, "maximum": 128, "step": 1}), + "sdnq_codebook_steps": OptionInfo(24, "Lloyd-Max Codebook steps", gr.Slider, {"minimum": 1, "maximum": 128, "step": 1}), "sdnq_dynamic_loss_threshold": OptionInfo(-1, "Dynamic loss threshold", gr.Slider, {"minimum": -1, "maximum": 0.1, "step": 1e-4}), "sdnq_use_svd": OptionInfo(False, "Use SVD quantization", gr.Checkbox), "sdnq_use_hadamard": OptionInfo(False, "Use Hadamard rotations", gr.Checkbox), + "sdnq_use_codebook": OptionInfo(False, "Use Lloyd-Max Codebook quantization", gr.Checkbox), "sdnq_use_dynamic_quantization": OptionInfo(False, "Use Dynamic quantization", gr.Checkbox), "sdnq_quantize_conv_layers": OptionInfo(False, "Quantize convolutional layers", gr.Checkbox), "sdnq_quantize_embedding_layers": OptionInfo(False, "Quantize embedding layers", gr.Checkbox),