Add loading image processor (#2254)

* add loading image processor

* fix style

* fix pix2struct

Author: eaidova

optimum/exporters/onnx/model_configs.py

@@ -2579,7 +2579,7 @@ def _create_dummy_input_generator_classes(self, **kwargs) -> List["DummyInputGen
         dummy_inputs_generators = []
         dummy_inputs_generators.append(self.DUMMY_INPUT_GENERATOR_CLASSES[0](self.task, self._normalized_config))
 
-        if self._preprocessors is None or len(self._preprocessors) != 2:
+        if self._preprocessors is None or len(self._preprocessors) < 2:
             raise ValueError(
                 f"Preprocessors for pix2struct need to be available for the ONNX export to infer input static shapes. Got: {self._preprocessors}"
             )
@@ -2598,7 +2598,7 @@ def _create_dummy_input_generator_classes(self, **kwargs) -> List["DummyInputGen
     def overwrite_shape_and_generate_input(
         self, dummy_input_gen: "DummyInputGenerator", input_name: str, framework: str, input_shapes: Dict
     ):
-        if self._preprocessors is None or len(self._preprocessors) != 2:
+        if self._preprocessors is None or len(self._preprocessors) < 2:
             raise ValueError(
                 f"Preprocessors for pix2struct need to be available for the ONNX export to infer input static shapes. Got: {self._preprocessors}"
             )

optimum/utils/save_utils.py

@@ -18,7 +18,7 @@
 from pathlib import Path
 from typing import List, Union
 
-from transformers import AutoFeatureExtractor, AutoProcessor, AutoTokenizer
+from transformers import AutoFeatureExtractor, AutoImageProcessor, AutoProcessor, AutoTokenizer
 
 
 logger = logging.getLogger(__name__)
@@ -50,6 +50,15 @@ def maybe_load_preprocessors(
         )
     except Exception:
         pass
+
+    try:
+        preprocessors.append(
+            AutoImageProcessor.from_pretrained(
+                src_name_or_path, subfolder=subfolder, trust_remote_code=trust_remote_code
+            )
+        )
+    except Exception:
+        pass
     return preprocessors