fix tests

Author: eaidova

optimum/exporters/openvino/convert.py

@@ -1008,7 +1008,11 @@ def get_diffusion_models_for_export_ext(
             task="feature-extraction",
             model_type="t5-encoder-model",
         )
-        export_config = export_config_constructor(text_encoder_3.config, int_dtype=int_dtype, float_dtype=float_dtype, )
+        export_config = export_config_constructor(
+            text_encoder_3.config,
+            int_dtype=int_dtype,
+            float_dtype=float_dtype,
+        )
         models_for_export["text_encoder_3"] = (text_encoder_3, export_config)
 
     return None, models_for_export

optimum/exporters/openvino/model_configs.py

@@ -20,7 +20,6 @@
 from packaging import version
 from transformers import PretrainedConfig, PreTrainedModel, TFPreTrainedModel
 from transformers.utils import is_tf_available
-from optimum.exporters.onnx.model_patcher import ModelPatcher
 
 from optimum.exporters.onnx.config import OnnxConfig, TextDecoderOnnxConfig, TextDecoderWithPositionIdsOnnxConfig
 from optimum.exporters.onnx.model_configs import (
@@ -38,6 +37,7 @@
     UNetOnnxConfig,
     VisionOnnxConfig,
 )
+from optimum.exporters.onnx.model_patcher import ModelPatcher
 from optimum.exporters.tasks import TasksManager
 from optimum.utils import DEFAULT_DUMMY_SHAPES
 from optimum.utils.input_generators import (
@@ -1583,4 +1583,3 @@ def patch_model_for_export(
         self, model: Union["PreTrainedModel", "TFPreTrainedModel"], model_kwargs: Optional[Dict[str, Any]] = None
     ) -> ModelPatcher:
         return ModelPatcher(self, model, model_kwargs=model_kwargs)
-

optimum/intel/openvino/modeling_diffusion.py

@@ -576,7 +576,7 @@ def width(self) -> int:
     @property
     def batch_size(self) -> int:
         model = self.unet.model if self.unet is not None else self.transformer.model
-        batch_size = model.inputs[0].get_partial_shape()[0] 
+        batch_size = model.inputs[0].get_partial_shape()[0]
         if batch_size.is_dynamic:
             return -1
         return batch_size.get_length()
@@ -642,7 +642,7 @@ def _reshape_transformer(
         else:
             # The factor of 2 comes from the guidance scale > 1
             batch_size *= 2 * num_images_per_prompt
-        
+
         height = height // self.vae_scale_factor if height > 0 else height
         width = width // self.vae_scale_factor if width > 0 else width
         shapes = {}
@@ -954,7 +954,9 @@ def modules(self):
 class OVModelTextEncoder(OVPipelinePart):
     def __init__(self, model: openvino.runtime.Model, parent_pipeline: OVDiffusionPipeline, model_name: str = ""):
         super().__init__(model, parent_pipeline, model_name)
-        self.hidden_states_output_names = sorted({name for out in self.model.outputs for name in out.names if name.startswith("hidden_states")})
+        self.hidden_states_output_names = sorted(
+            {name for out in self.model.outputs for name in out.names if name.startswith("hidden_states")}
+        )
 
     def forward(
         self,
@@ -971,9 +973,13 @@ def forward(
         main_out = ov_outputs[0]
         model_outputs = {}
         model_outputs[self.model.outputs[0].get_any_name()] = torch.from_numpy(main_out)
-        if self.hidden_states_output_names and not "last_hidden_state" in model_outputs:
+        if self.hidden_states_output_names and "last_hidden_state" not in model_outputs:
             model_outputs["last_hidden_state"] = torch.from_numpy(ov_outputs[self.hidden_states_output_names[-1]])
-        if self.hidden_states_output_names and output_hidden_states or self.config.output_hidden_states:
+        if (
+            self.hidden_states_output_names
+            and output_hidden_states
+            or getattr(self.config, "output_hidden_states", False)
+        ):
             hidden_states = [torch.from_numpy(ov_outputs[out_name]) for out_name in self.hidden_states_output_names]
             model_outputs["hidden_states"] = hidden_states
 

tests/openvino/test_diffusion.py

@@ -185,10 +185,19 @@ def test_shape(self, model_arch: str):
             elif output_type == "pt":
                 self.assertEqual(outputs.shape, (batch_size, 3, height, width))
             else:
-                out_channels = pipeline.unet.config.out_channels if pipeline.unet is not None else pipeline.transformer.config.out_channels
+                out_channels = (
+                    pipeline.unet.config.out_channels
+                    if pipeline.unet is not None
+                    else pipeline.transformer.config.out_channels
+                )
                 self.assertEqual(
                     outputs.shape,
-                    (batch_size, out_channels, height // pipeline.vae_scale_factor, width // pipeline.vae_scale_factor),
+                    (
+                        batch_size,
+                        out_channels,
+                        height // pipeline.vae_scale_factor,
+                        width // pipeline.vae_scale_factor,
+                    ),
                 )
 
     @parameterized.expand(SUPPORTED_ARCHITECTURES)
@@ -246,7 +255,7 @@ def test_negative_prompt(self, model_arch: str):
                 do_classifier_free_guidance=True,
                 negative_prompt=negative_prompt,
             )
-    
+
         else:
             inputs["prompt_embeds"], inputs["negative_prompt_embeds"] = pipeline.encode_prompt(
                 prompt=prompt,
@@ -306,8 +315,10 @@ def test_height_width_properties(self, model_arch: str):
         )
 
         self.assertFalse(ov_pipeline.is_dynamic)
-        expected_batch  =  batch_size * num_images_per_prompt
-        if ov_pipeline.unet is not None and "timestep_cond" not in {inputs.get_any_name() for inputs in ov_pipeline.unet.model.inputs}:
+        expected_batch = batch_size * num_images_per_prompt
+        if ov_pipeline.unet is not None and "timestep_cond" not in {
+            inputs.get_any_name() for inputs in ov_pipeline.unet.model.inputs
+        }:
             expected_batch *= 2
         self.assertEqual(
             ov_pipeline.batch_size,
@@ -435,10 +446,19 @@ def test_shape(self, model_arch: str):
                 elif output_type == "pt":
                     self.assertEqual(outputs.shape, (batch_size, 3, height, width))
                 else:
-                    out_channels = pipeline.unet.config.out_channels if pipeline.unet is not None else pipeline.transformer.config.out_channels
+                    out_channels = (
+                        pipeline.unet.config.out_channels
+                        if pipeline.unet is not None
+                        else pipeline.transformer.config.out_channels
+                    )
                     self.assertEqual(
                         outputs.shape,
-                        (batch_size, out_channels, height // pipeline.vae_scale_factor, width // pipeline.vae_scale_factor),
+                        (
+                            batch_size,
+                            out_channels,
+                            height // pipeline.vae_scale_factor,
+                            width // pipeline.vae_scale_factor,
+                        ),
                     )
 
     @parameterized.expand(SUPPORTED_ARCHITECTURES)
@@ -521,13 +541,12 @@ def test_height_width_properties(self, model_arch: str):
         )
 
         self.assertFalse(ov_pipeline.is_dynamic)
-        expected_batch  =  batch_size * num_images_per_prompt
-        if ov_pipeline.unet is not None and "timestep_cond" not in {inputs.get_any_name() for inputs in ov_pipeline.unet.model.inputs}:
+        expected_batch = batch_size * num_images_per_prompt
+        if ov_pipeline.unet is not None and "timestep_cond" not in {
+            inputs.get_any_name() for inputs in ov_pipeline.unet.model.inputs
+        }:
             expected_batch *= 2
-        self.assertEqual(
-            ov_pipeline.batch_size,
-            expected_batch
-        )
+        self.assertEqual(ov_pipeline.batch_size, expected_batch)
         self.assertEqual(ov_pipeline.height, height)
         self.assertEqual(ov_pipeline.width, width)
 
@@ -655,10 +674,19 @@ def test_shape(self, model_arch: str):
                 elif output_type == "pt":
                     self.assertEqual(outputs.shape, (batch_size, 3, height, width))
                 else:
-                    out_channels = pipeline.unet.config.out_channels if pipeline.unet is not None else pipeline.transformer.config.out_channels
+                    out_channels = (
+                        pipeline.unet.config.out_channels
+                        if pipeline.unet is not None
+                        else pipeline.transformer.config.out_channels
+                    )
                     self.assertEqual(
                         outputs.shape,
-                        (batch_size, out_channels, height // pipeline.vae_scale_factor, width // pipeline.vae_scale_factor),
+                        (
+                            batch_size,
+                            out_channels,
+                            height // pipeline.vae_scale_factor,
+                            width // pipeline.vae_scale_factor,
+                        ),
                     )
 
     @parameterized.expand(SUPPORTED_ARCHITECTURES)
@@ -676,7 +704,7 @@ def test_compare_to_diffusers_pipeline(self, model_arch: str):
 
             ov_output = ov_pipeline(**inputs, generator=get_generator("pt", SEED)).images
             diffusers_output = diffusers_pipeline(**inputs, generator=get_generator("pt", SEED)).images
-            
+
             np.testing.assert_allclose(ov_output, diffusers_output, atol=6e-3, rtol=1e-2)
 
     @parameterized.expand(SUPPORTED_ARCHITECTURES)
@@ -741,8 +769,10 @@ def test_height_width_properties(self, model_arch: str):
         )
 
         self.assertFalse(ov_pipeline.is_dynamic)
-        expected_batch  =  batch_size * num_images_per_prompt
-        if ov_pipeline.unet is not None and "timestep_cond" not in {inputs.get_any_name() for inputs in ov_pipeline.unet.model.inputs}:
+        expected_batch = batch_size * num_images_per_prompt
+        if ov_pipeline.unet is not None and "timestep_cond" not in {
+            inputs.get_any_name() for inputs in ov_pipeline.unet.model.inputs
+        }:
             expected_batch *= 2
         self.assertEqual(
             ov_pipeline.batch_size,