Quantize SD submodels in OVQuantizer

Author: nikita-savelyevv

notebooks/openvino/stable_diffusion_hybrid_quantization.ipynb

@@ -53,6 +53,7 @@
     "from pathlib import Path\n",
     "from openvino.runtime import Core\n",
     "from optimum.intel import OVConfig, OVQuantizer, OVStableDiffusionPipeline, OVWeightQuantizationConfig\n",
+    "from optimum.intel.openvino.configuration import OVQuantizationMethod\n",
     "\n",
     "transformers.logging.set_verbosity_error()\n",
     "datasets.logging.set_verbosity_error()"
@@ -198,8 +199,8 @@
    },
    "outputs": [],
    "source": [
-    "quantization_config = OVWeightQuantizationConfig(bits=8, num_samples=NUM_SAMPLES)\n",
     "int8_pipe = OVStableDiffusionPipeline.from_pretrained(model_id=MODEL_ID, export=True)\n",
+    "quantization_config = OVWeightQuantizationConfig(bits=8, num_samples=NUM_SAMPLES, quant_method=OVQuantizationMethod.HYBRID)\n",
     "quantizer = OVQuantizer(int8_pipe)\n",
     "quantizer.quantize(\n",
     "    ov_config=OVConfig(quantization_config=quantization_config),\n",
@@ -618,7 +619,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.8.10"
+   "version": "3.11.7"
   }
  },
  "nbformat": 4,

optimum/intel/openvino/modeling_diffusion.py

@@ -292,33 +292,27 @@ def _from_pretrained(
                 else:
                     kwargs[name] = load_method(new_model_save_dir)
 
-        quantization_config = cls._prepare_weight_quantization_config(quantization_config, load_in_8bit)
-
         unet_path = new_model_save_dir / DIFFUSION_MODEL_UNET_SUBFOLDER / unet_file_name
-        if quantization_config is not None and quantization_config.dataset is not None:
-            # load the UNet model uncompressed to apply hybrid quantization further
-            unet = cls.load_model(unet_path)
-            # Apply weights compression to other `components` without dataset
-            quantization_config_without_dataset = deepcopy(quantization_config)
-            quantization_config_without_dataset.dataset = None
-        else:
-            quantization_config_without_dataset = quantization_config
-            unet = cls.load_model(unet_path, quantization_config_without_dataset)
-
         components = {
             "vae_encoder": new_model_save_dir / DIFFUSION_MODEL_VAE_ENCODER_SUBFOLDER / vae_encoder_file_name,
             "vae_decoder": new_model_save_dir / DIFFUSION_MODEL_VAE_DECODER_SUBFOLDER / vae_decoder_file_name,
             "text_encoder": new_model_save_dir / DIFFUSION_MODEL_TEXT_ENCODER_SUBFOLDER / text_encoder_file_name,
             "text_encoder_2": new_model_save_dir / DIFFUSION_MODEL_TEXT_ENCODER_2_SUBFOLDER / text_encoder_2_file_name,
         }
 
-        for key, value in components.items():
-            components[key] = cls.load_model(value, quantization_config_without_dataset) if value.is_file() else None
-
         if model_save_dir is None:
             model_save_dir = new_model_save_dir
 
-        if quantization_config is not None and quantization_config.dataset is not None:
+        quantization_config = cls._prepare_weight_quantization_config(quantization_config, load_in_8bit)
+        if quantization_config is None or quantization_config.dataset is None:
+            unet = cls.load_model(unet_path, quantization_config)
+            for key, value in components.items():
+                components[key] = cls.load_model(value, quantization_config) if value.is_file() else None
+        else:
+            # Load uncompressed models to apply hybrid quantization further
+            unet = cls.load_model(unet_path)
+            for key, value in components.items():
+                components[key] = cls.load_model(value) if value.is_file() else None
             sd_model = cls(unet=unet, config=config, model_save_dir=model_save_dir, **components, **kwargs)
 
             supported_pipelines = (
@@ -331,10 +325,10 @@ def _from_pretrained(
 
             from optimum.intel import OVQuantizer
 
+            hybrid_quantization_config = deepcopy(quantization_config)
+            hybrid_quantization_config.quant_method = OVQuantizationMethod.HYBRID
             quantizer = OVQuantizer(sd_model)
-            quantization_config_copy = deepcopy(quantization_config)
-            quantization_config_copy.quant_method = OVQuantizationMethod.HYBRID
-            quantizer.quantize(ov_config=OVConfig(quantization_config=quantization_config_copy))
+            quantizer.quantize(ov_config=OVConfig(quantization_config=hybrid_quantization_config))
 
             return sd_model
 
@@ -347,6 +341,7 @@ def _from_pretrained(
             **kwargs,
         )
 
+
     @classmethod
     def _from_transformers(
         cls,

optimum/intel/openvino/quantization.py

@@ -63,8 +63,6 @@
 
 if is_datasets_available():
     from datasets import Dataset
-else:
-    Dataset = None
 
 register_module(ignored_algorithms=[])(Conv1D)
 
@@ -388,11 +386,21 @@ def _quantize_ovbasemodel(
                 if calibration_dataset is None:
                     raise ValueError("Calibration dataset is required to run hybrid quantization.")
                 if isinstance(self.model, OVStableDiffusionPipelineBase):
+                    # Apply weight-only quantization to all SD submodels except UNet
+                    quantization_config_copy = copy.deepcopy(quantization_config)
+                    quantization_config_copy.dataset = None
+                    quantization_config_copy.quant_method = OVQuantizationMethod.DEFAULT
+                    for sd_submodel_name in ["vae_encoder", "vae_decoder", "text_encoder", "text_encoder_2"]:
+                        sd_submodel = getattr(self.model, sd_submodel_name)
+                        if sd_submodel is not None:
+                            _weight_only_quantization(sd_submodel.model, quantization_config_copy)
+
+                    # Apply hybrid quantization to UNet
                     self.model.unet.model = _hybrid_quantization(
                         self.model.unet.model, quantization_config, calibration_dataset
                     )
                 else:
-                    # This may be for example OVModelForImageClassification, OVModelForAudioClassification, etc.
+                    # The model may be for example OVModelForImageClassification, OVModelForAudioClassification, etc.
                     self.model.model = _hybrid_quantization(self.model.model, quantization_config, calibration_dataset)
             else:
                 _weight_only_quantization(self.model.model, quantization_config, calibration_dataset)