diff --git a/optimum/exporters/openvino/convert.py b/optimum/exporters/openvino/convert.py index 5b9aea10e2..e4ece9801b 100644 --- a/optimum/exporters/openvino/convert.py +++ b/optimum/exporters/openvino/convert.py @@ -98,10 +98,11 @@ def _set_runtime_options( Tuple[Union["PreTrainedModel", "TFPreTrainedModel", "ModelMixin", "DiffusionPipeline"], "OnnxConfig"], ], task: str, + library_name: str, ): for model_name in models_and_export_configs.keys(): _, sub_export_config = models_and_export_configs[model_name] - if "vae_" in model_name or "text-generation" in task: + if "diffusers" in library_name or "text-generation" in task: sub_export_config.runtime_options = {"ACTIVATIONS_SCALE_FACTOR": "8.0"} @@ -754,7 +755,7 @@ def export_from_model( model.save_config(output) - _set_runtime_options(models_and_export_configs, task) + _set_runtime_options(models_and_export_configs, task, library_name) export_models( models_and_export_configs=models_and_export_configs, diff --git a/tests/openvino/test_export.py b/tests/openvino/test_export.py index 4c42f8a337..80a45cab6e 100644 --- a/tests/openvino/test_export.py +++ b/tests/openvino/test_export.py @@ -140,6 +140,26 @@ def _openvino_export( self.assertTrue( ov_model.vae_decoder.model.has_rt_info(["runtime_options", "ACTIVATIONS_SCALE_FACTOR"]) ) + if hasattr(ov_model, "text_encoder") and ov_model.text_encoder: + self.assertTrue( + ov_model.text_encoder.model.has_rt_info(["runtime_options", "ACTIVATIONS_SCALE_FACTOR"]) + ) + if hasattr(ov_model, "text_encoder_2") and ov_model.text_encoder_2: + self.assertTrue( + ov_model.text_encoder_2.model.has_rt_info(["runtime_options", "ACTIVATIONS_SCALE_FACTOR"]) + ) + if hasattr(ov_model, "text_encoder_3") and ov_model.text_encoder_3: + self.assertTrue( + ov_model.text_encoder_3.model.has_rt_info(["runtime_options", "ACTIVATIONS_SCALE_FACTOR"]) + ) + if hasattr(ov_model, "unet") and ov_model.unet: + self.assertTrue( + ov_model.unet.model.has_rt_info(["runtime_options", "ACTIVATIONS_SCALE_FACTOR"]) + ) + if hasattr(ov_model, "transformer") and ov_model.transformer: + self.assertTrue( + ov_model.transformer.model.has_rt_info(["runtime_options", "ACTIVATIONS_SCALE_FACTOR"]) + ) @parameterized.expand(SUPPORTED_ARCHITECTURES) def test_export(self, model_type: str):