Add Lloyd-Max quantization with use_codebook

2026-09-19 01:04:32 +02:00 · 2026-08-10 22:51:22 +03:00
parent ccb048a9af
commit 251945939b
3 changed files with 11 additions and 2 deletions
@@ -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,
@@ -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


@@ -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),