diff --git a/extensions-builtin/Lora/network_lora.py b/extensions-builtin/Lora/network_lora.py index b31536e7d..a9e78b6e4 100644 --- a/extensions-builtin/Lora/network_lora.py +++ b/extensions-builtin/Lora/network_lora.py @@ -25,8 +25,8 @@ class NetworkModuleLora(network.NetworkModule): if weight is None and none_ok: return None linear_modules = [torch.nn.Linear, torch.nn.modules.linear.NonDynamicallyQuantizableLinear, torch.nn.MultiheadAttention, diffusers_lora.LoRACompatibleLinear] - is_linear = type(self.sd_module) in linear_modules or self.sd_module.__class__.__name__ == "NNCFLinear" - is_conv = type(self.sd_module) in [torch.nn.Conv2d, diffusers_lora.LoRACompatibleConv] or self.sd_module.__class__.__name__ == "NNCFConv2d" + is_linear = type(self.sd_module) in linear_modules or self.sd_module.__class__.__name__ in {"NNCFLinear", "QLinear"} + is_conv = type(self.sd_module) in [torch.nn.Conv2d, diffusers_lora.LoRACompatibleConv] or self.sd_module.__class__.__name__ in {"NNCFConv2d", "QConv2d"} if is_linear: weight = weight.reshape(weight.shape[0], -1) module = torch.nn.Linear(weight.shape[1], weight.shape[0], bias=False) diff --git a/installer.py b/installer.py index 6201c17bb..8fda1e86b 100644 --- a/installer.py +++ b/installer.py @@ -626,6 +626,8 @@ def install_torch_addons(): install('olive-ai') if opts.get('nncf_compress_weights', False) and not args.use_openvino: install('nncf==2.7.0', 'nncf') + if opts.get('optimum_quanto_weights', False): + install('optimum-quanto', 'optimum-quanto') if triton_command is not None: install(triton_command, 'triton', quiet=True) diff --git a/modules/control/units/controlnet.py b/modules/control/units/controlnet.py index ec99b7bd6..40f543d1b 100644 --- a/modules/control/units/controlnet.py +++ b/modules/control/units/controlnet.py @@ -189,6 +189,15 @@ class ControlNet(): self.model = nncf_compress_model(self.model) except Exception as e: log.error(f'Control {what} model NNCF Compression failed: id="{model_id}" error={e}') + elif "ControlNet" in opts.optimum_quanto_weights: + try: + log.debug(f'Control {what} model Optimum Quanto: id="{model_id}"') + from installer import install + install('optimum-quanto', quiet=True) + from modules.sd_models_compile import optimum_quanto_model + self.model = optimum_quanto_model(self.model) + except Exception as e: + log.error(f'Control {what} model Optimum Quanto failed: id="{model_id}" error={e}') if self.device is not None: self.model.to(self.device) t1 = time.time() diff --git a/modules/sd_models.py b/modules/sd_models.py index 9109b6a7f..f215dffe1 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -1180,6 +1180,8 @@ def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=No set_diffuser_options(sd_model, vae, op, offload=False) if shared.opts.nncf_compress_weights and not (shared.opts.cuda_compile and shared.opts.cuda_compile_backend == "openvino_fx"): sd_model = sd_models_compile.nncf_compress_weights(sd_model) # run this before move model so it can be compressed in CPU + if shared.opts.optimum_quanto_weights: + sd_model = sd_models_compile.optimum_quanto_weights(sd_model) # run this before move model so it can be compressed in CPU timer.record("options") set_diffuser_offload(sd_model, op) diff --git a/modules/sd_models_compile.py b/modules/sd_models_compile.py index 4d4776c26..bf6e61804 100644 --- a/modules/sd_models_compile.py +++ b/modules/sd_models_compile.py @@ -145,11 +145,37 @@ def nncf_compress_weights(sd_model): sd_model = apply_compile_to_model(sd_model, nncf_compress_model, shared.opts.nncf_compress_weights, op="nncf") t1 = time.time() - shared.log.info(f"Compress Weights: time={t1-t0:.2f}") + shared.log.info(f"NNCF Compress Weights: time={t1-t0:.2f}") except Exception as e: - shared.log.warning(f"Compress Weights: error: {e}") + shared.log.warning(f"NNCF Compress Weights: error: {e}") return sd_model +def optimum_quanto_model(model): + from optimum import quanto + model.eval() + backup_embeddings = None + if hasattr(model, "get_input_embeddings"): + backup_embeddings = copy.deepcopy(model.get_input_embeddings()) + quanto.quantize(model, weights=getattr(quanto, shared.opts.optimum_quanto_weights_type)) + quanto.freeze(model) + if hasattr(model, "set_input_embeddings") and backup_embeddings is not None: + model.set_input_embeddings(backup_embeddings) + devices.torch_gc(force=True) + return model + +def optimum_quanto_weights(sd_model): + try: + t0 = time.time() + from installer import install + install('optimum-quanto', quiet=True) + + sd_model = apply_compile_to_model(sd_model, optimum_quanto_model, shared.opts.optimum_quanto_weights, op="optimum-quanto") + + t1 = time.time() + shared.log.info(f"Optimum Quanto Weights: time={t1-t0:.2f}") + except Exception as e: + shared.log.warning(f"Optimum Quanto Weights: error: {e}") + return sd_model def optimize_openvino(sd_model): try: diff --git a/modules/shared.py b/modules/shared.py index 1cf31c491..3b4a680ff 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -446,11 +446,13 @@ options_templates.update(options_section(('cuda', "Compute Settings"), { "cuda_compile_precompile": OptionInfo(False, "Model compile precompile"), "cuda_compile_verbose": OptionInfo(False, "Model compile verbose mode"), "cuda_compile_errors": OptionInfo(True, "Model compile suppress errors"), - "diffusers_quantization": OptionInfo(False, "Dynamic quantization with TorchAO"), "deep_cache_interval": OptionInfo(3, "DeepCache cache interval", gr.Slider, {"minimum": 1, "maximum": 10, "step": 1}), - "nncf_sep": OptionInfo("

Model Compress

", "", gr.HTML), - "nncf_compress_weights": OptionInfo([], "Compress Model weights with NNCF", gr.CheckboxGroup, {"choices": ["Model", "VAE", "Text Encoder", "ControlNet"], "visible": native}), + "quant_sep": OptionInfo("

Model Quantization

", "", gr.HTML), + "diffusers_quantization": OptionInfo(False, "Dynamic quantization with TorchAO"), + "nncf_compress_weights": OptionInfo([], "Compress Model weights with NNCF INT8", gr.CheckboxGroup, {"choices": ["Model", "VAE", "Text Encoder", "ControlNet"], "visible": native}), + "optimum_quanto_weights": OptionInfo([], "Quantize Model weights with Optimum Quanto", gr.CheckboxGroup, {"choices": ["Model", "VAE", "Text Encoder", "ControlNet"], "visible": native}), + "optimum_quanto_weights_type": OptionInfo("qint8", "Quant mode for Optimum Quanto", gr.Radio, {"choices": ['qint8', 'qfloat8_e4m3fn', 'qfloat8_e5m2', 'qint4', 'qint2'], "visible": native}), "ipex_sep": OptionInfo("

IPEX

", "", gr.HTML, {"visible": devices.backend == "ipex"}), "ipex_optimize": OptionInfo([], "IPEX Optimize for Intel GPUs", gr.CheckboxGroup, {"choices": ["Model", "VAE", "Text Encoder", "Upscaler"], "visible": devices.backend == "ipex"}),