Expand activation scaling to other submodels of diffusion pipelines (#…

…1039) * Expand activation scaling to other submodels * Applied comments
huggingface · Nov 27, 2024 · c454b00 · c454b00
1 parent f6a7b83
commit c454b00
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 2 deletions.
diff --git 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
@@ -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):