Optimum Quanto support

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)

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)
 

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()

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)

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:

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("<h2>Model Compress</h2>", "", gr.HTML),
-    "nncf_compress_weights": OptionInfo([], "Compress Model weights with NNCF", gr.CheckboxGroup, {"choices": ["Model", "VAE", "Text Encoder", "ControlNet"], "visible": native}),
+    "quant_sep": OptionInfo("<h2>Model Quantization</h2>", "", 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("<h2>IPEX</h2>", "", 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"}),