fix import

huggingface · Oct 9, 2024 · 354abdf · 354abdf
1 parent c540c7b
commit 354abdf
Showing 1 changed file with 5 additions and 5 deletions.
diff --git a/optimum/intel/openvino/quantization.py b/optimum/intel/openvino/quantization.py
@@ -314,7 +314,7 @@ def _quantize_ovbasemodel(
         **kwargs,
     ):
         if is_diffusers_available():
-            from optimum.intel.openvino.modeling_diffusion import OVPipeline
+            from optimum.intel.openvino.modeling_diffusion import OVDiffusionPipeline
 
         if save_directory is not None:
             save_directory = Path(save_directory)
@@ -324,7 +324,7 @@ def _quantize_ovbasemodel(
         if calibration_dataset is not None:
             # Process custom calibration dataset
 
-            if is_diffusers_available() and isinstance(self.model, OVPipeline):
+            if is_diffusers_available() and isinstance(self.model, OVDiffusionPipeline):
                 calibration_dataset = self._prepare_unet_dataset(
                     quantization_config.num_samples, dataset=calibration_dataset
                 )
@@ -361,7 +361,7 @@ def _quantize_ovbasemodel(
 
                 if isinstance(self.model, OVModelForCausalLM):
                     calibration_dataset = self._prepare_causal_lm_dataset(quantization_config)
-                elif is_diffusers_available() and isinstance(self.model, OVPipeline):
+                elif is_diffusers_available() and isinstance(self.model, OVDiffusionPipeline):
                     if not isinstance(quantization_config.dataset, str):
                         raise ValueError("Please provide dataset as one of the accepted dataset labels.")
                     calibration_dataset = self._prepare_unet_dataset(
@@ -375,7 +375,7 @@ def _quantize_ovbasemodel(
             if quantization_config.quant_method == OVQuantizationMethod.HYBRID:
                 if calibration_dataset is None:
                     raise ValueError("Calibration dataset is required to run hybrid quantization.")
-                if is_diffusers_available() and isinstance(self.model, OVPipeline):
+                if is_diffusers_available() and isinstance(self.model, OVDiffusionPipeline):
                     # Apply weight-only quantization to all SD submodels except UNet
                     quantization_config_copy = copy.deepcopy(quantization_config)
                     quantization_config_copy.dataset = None
@@ -395,7 +395,7 @@ def _quantize_ovbasemodel(
                     self.model.model = _hybrid_quantization(self.model.model, quantization_config, calibration_dataset)
                     self.model.request = None
             else:
-                if is_diffusers_available() and isinstance(self.model, OVPipeline):
+                if is_diffusers_available() and isinstance(self.model, OVDiffusionPipeline):
                     sub_model_names = ["vae_encoder", "vae_decoder", "text_encoder", "text_encoder_2", "unet"]
                     sub_models = filter(lambda x: x, (getattr(self.model, name) for name in sub_model_names))
                     for sub_model in sub_models: