OpenVINO 4 bit support for CPUs

modules/intel/openvino/__init__.py

@@ -229,12 +229,16 @@ def openvino_compile(gm: GraphModule, *args, model_hash_str: str = None, file_na
         om.inputs[idx].get_node().set_element_type(dtype_mapping[input_data.dtype])
         om.inputs[idx].get_node().set_partial_shape(PartialShape(list(input_data.shape)))
     om.validate_nodes_and_infer_types()
-    if shared.opts.nncf_compress_weights:
-        om = nncf.compress_weights(om)
+    if shared.opts.nncf_compress_weights and not (shared.compiled_model_state.compiling_vae and not shared.opts.nncf_compress_vae_weights):
+        if shared.compiled_model_state.compiling_vae or shared.opts.nncf_compress_weights_mode == "INT8":
+            om = nncf.compress_weights(om)
+        else:
+            om = nncf.compress_weights(om, mode=getattr(nncf.CompressWeightsMode, shared.opts.nncf_compress_weights_mode), group_size=8, ratio=shared.opts.nncf_compress_weights_raito)
 
     if model_hash_str is not None:
         core.set_property({'CACHE_DIR': cache_root + '/blob'})
 
+    shared.compiled_model_state.compiling_vae = False
     compiled_model = core.compile_model(om, device)
     return compiled_model
 
@@ -257,13 +261,16 @@ def openvino_compile_cached_model(cached_model_path, *example_inputs):
         om.inputs[idx].get_node().set_element_type(dtype_mapping[input_data.dtype])
         om.inputs[idx].get_node().set_partial_shape(PartialShape(list(input_data.shape)))
     om.validate_nodes_and_infer_types()
-    if shared.opts.nncf_compress_weights:
-        om = nncf.compress_weights(om)
+    if shared.opts.nncf_compress_weights and not (shared.compiled_model_state.compiling_vae and not shared.opts.nncf_compress_vae_weights):
+        if shared.compiled_model_state.compiling_vae or shared.opts.nncf_compress_weights_mode == "INT8":
+            om = nncf.compress_weights(om)
+        else:
+            om = nncf.compress_weights(om, mode=getattr(nncf.CompressWeightsMode, shared.opts.nncf_compress_weights_mode), group_size=8, ratio=shared.opts.nncf_compress_weights_raito)
 
     core.set_property({'CACHE_DIR': shared.opts.openvino_cache_path + '/blob'})
 
+    shared.compiled_model_state.compiling_vae = False
     compiled_model = core.compile_model(om, get_device())
-
     return compiled_model
 
 def openvino_execute(gm: GraphModule, *args, executor_parameters=None, partition_id, file_name=""):

modules/processing_vae.py

@@ -48,6 +48,10 @@ def full_vae_decode(latents, model):
         model.upcast_vae()
         latents = latents.to(next(iter(model.vae.post_quant_conv.parameters())).dtype)
 
+    # OpenVINO with INT4 doesn't work with VAE decode so we pass that we are using VAE right now to OpenVINO
+    if shared.compiled_model_state is not None:
+        shared.compiled_model_state.compiling_vae = True
+
     decoded = model.vae.decode(latents / model.vae.config.scaling_factor, return_dict=False)[0]
 
     # Downcast VAE after OpenVINO compile

modules/sd_models_compile.py

@@ -20,6 +20,7 @@ class CompiledModelState:
         self.lora_compile = False
         self.compiled_cache = {}
         self.partitioned_modules = {}
+        self.compiling_vae = False
 
 
 def ipex_optimize(sd_model):

modules/shared.py

@@ -331,7 +331,9 @@ 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, "Enable dynamic quantization"),
+    "diffusers_quantization": OptionInfo(False, "Enable dynamic quantization with torchao"),
+    "nncf_compress_weights": OptionInfo(False, "Compress Model weights with NNCF"),
+    "nncf_compress_vae_weights": OptionInfo(False, "Compress VAE weights with NNCF"),
 
     "ipex_sep": OptionInfo("<h2>IPEX, DirectML and OpenVINO</h2>", "", gr.HTML),
     "ipex_optimize": OptionInfo(False if not devices.backend == "ipex" else True, "Enable IPEX Optimize for Intel GPUs"),
@@ -342,8 +344,8 @@ options_templates.update(options_section(('cuda', "Compute Settings"), {
     "openvino_hetero_gpu": OptionInfo(False, "OpenVINO use Hetero Device for single inference with multiple devices"),
     "openvino_remove_cpu_from_hetero": OptionInfo(False, "OpenVINO remove CPU from Hetero Device"),
     "openvino_remove_igpu_from_hetero": OptionInfo(False, "OpenVINO remove iGPU from Hetero Device"),
-    "nncf_compress_weights": OptionInfo(False, "Compress Model weights to 8 bit with NNCF"),
-    "nncf_compress_vae_weights": OptionInfo(False, "Compress VAE weights to 8 bit with NNCF"),
+    "nncf_compress_weights_mode": OptionInfo("INT8", "OpenVINO compress mode for NNCF (CPU Only)", gr.Radio, {"choices": ['INT8', 'INT4_SYM', 'INT4_ASYM', 'NF4']}),
+    "nncf_compress_weights_raito": OptionInfo(1.0, "OpenVINO compress ratio for NNCF with 4-bit modes", gr.Slider, {"minimum": 0, "maximum": 1, "step": 0.01}),
 }))
 
 options_templates.update(options_section(('advanced', "Inference Settings"), {