Repository: huggingface/optimum
Pull request #2255: Add feature extraction support for image models
Author: login: Abdennacer-Badaoui
Created: 2025-05-09T10:36:31Z | Merged: 2025-05-12T08:21:47Z

# What does this PR do?

This PR adds support for image models in ORTModelForFeatureExtraction, specifically addressing ViT and other visual models. The implementation modifies the forward method to dynamically handle both text and image inputs by detecting the presence of image-specific input keys like 'pixel_values'. This allows users to seamlessly use ORTModelForFeatureExtraction for both text and image feature extraction tasks without requiring a separate class structure. Tests have been updated to properly validate this functionality across all supported architectures, including Vision Transformer (ViT) models.


Fixes # (issue)
Fixes the issue : https://github.com/huggingface/optimum/issues/2226

## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [ ] Did you make sure to update the documentation with your changes?
- [x] Did you write any new necessary tests?




Diff:
diff --git a/optimum/onnxruntime/modeling_ort.py b/optimum/onnxruntime/modeling_ort.py
index 8d783a4b39..41ac8d4c5f 100644
--- a/optimum/onnxruntime/modeling_ort.py
+++ b/optimum/onnxruntime/modeling_ort.py
@@ -1059,18 +1059,37 @@ def forward(
         input_ids: Optional[Union[torch.Tensor, np.ndarray]] = None,
         attention_mask: Optional[Union[torch.Tensor, np.ndarray]] = None,
         token_type_ids: Optional[Union[torch.Tensor, np.ndarray]] = None,
-        **kwargs,
+        position_ids: Optional[Union[torch.Tensor, np.ndarray]] = None,
+        pixel_values: Optional[Union[torch.Tensor, np.ndarray]] = None,
+        input_features: Optional[Union[torch.Tensor, np.ndarray]] = None,
+        input_values: Optional[Union[torch.Tensor, np.ndarray]] = None,
     ):
-        use_torch = isinstance(input_ids, torch.Tensor)
+        # Determine the tensor type from any available tensor input
+        tensor_inputs = [
+            input_ids,
+            attention_mask,
+            token_type_ids,
+            position_ids,
+            pixel_values,
+            input_features,
+            input_values,
+        ]
+        first_tensor = next(filter(lambda x: x is not None, tensor_inputs))
+        use_torch = isinstance(first_tensor, torch.Tensor)
         self.raise_on_numpy_input_io_binding(use_torch)
 
         if token_type_ids is None and "token_type_ids" in self.input_names:
             token_type_ids = torch.zeros_like(input_ids) if use_torch else np.zeros_like(input_ids)
 
+        # Build model_inputs dictionary
         model_inputs = {
             "input_ids": input_ids,
             "attention_mask": attention_mask,
             "token_type_ids": token_type_ids,
+            "position_ids": position_ids,
+            "pixel_values": pixel_values,
+            "input_features": input_features,
+            "input_values": input_values,
         }
 
         if self.use_io_binding:
diff --git a/tests/onnxruntime/test_modeling.py b/tests/onnxruntime/test_modeling.py
index 52774c4563..aae2fa50d8 100644
--- a/tests/onnxruntime/test_modeling.py
+++ b/tests/onnxruntime/test_modeling.py
@@ -2280,6 +2280,141 @@ def test_default_token_type_ids(self):
         gc.collect()
 
 
+class ORTModelForFeatureExtractionFromImageModelsIntegrationTest(ORTModelTestMixin):
+    SUPPORTED_ARCHITECTURES = ["vit", "dinov2"]
+
+    FULL_GRID = {"model_arch": SUPPORTED_ARCHITECTURES}
+    ORTMODEL_CLASS = ORTModelForFeatureExtraction
+    TASK = "feature-extraction"
+
+    def get_raw_input(self, model_arch):
+        image_url = "https://picsum.photos/id/237/200/300"
+        return Image.open(requests.get(image_url, stream=True).raw)
+
+    def get_input(self, model_arch, processor, return_tensors="pt"):
+        raw_input = self.get_raw_input(model_arch)
+        return processor(images=raw_input, return_tensors=return_tensors)
+
+    @parameterized.expand(SUPPORTED_ARCHITECTURES)
+    def test_compare_to_transformers(self, model_arch):
+        model_args = {"test_name": model_arch, "model_arch": model_arch}
+        self._setup(model_args)
+
+        model_id = MODEL_NAMES[model_arch]
+        onnx_model = ORTModelForFeatureExtraction.from_pretrained(self.onnx_model_dirs[model_arch])
+
+        self.assertIsInstance(onnx_model.model, onnxruntime.InferenceSession)
+        self.assertIsInstance(onnx_model.config, PretrainedConfig)
+
+        set_seed(SEED)
+        transformers_model = AutoModel.from_pretrained(model_id)
+        processor = get_preprocessor(model_id)
+        inputs = self.get_input(model_arch, processor, return_tensors="pt")
+        with torch.no_grad():
+            transformers_outputs = transformers_model(**inputs)
+
+        for input_type in ["pt", "np"]:
+            inputs = self.get_input(model_arch, processor, return_tensors=input_type)
+            onnx_outputs = onnx_model(**inputs)
+
+            self.assertIn("last_hidden_state", onnx_outputs)
+            self.assertIsInstance(onnx_outputs.last_hidden_state, self.TENSOR_ALIAS_TO_TYPE[input_type])
+
+            # compare tensor outputs
+            torch.testing.assert_close(
+                torch.Tensor(onnx_outputs.last_hidden_state),
+                transformers_outputs.last_hidden_state,
+                atol=self.ATOL,
+                rtol=self.RTOL,
+            )
+
+        gc.collect()
+
+    @parameterized.expand(SUPPORTED_ARCHITECTURES)
+    def test_ort_model_inference(self, model_arch):
+        model_args = {"test_name": model_arch, "model_arch": model_arch}
+        self._setup(model_args)
+
+        model_id = MODEL_NAMES[model_arch]
+        onnx_model = ORTModelForFeatureExtraction.from_pretrained(self.onnx_model_dirs[model_arch])
+        processor = get_preprocessor(model_id)
+        raw_input = self.get_raw_input(model_arch)
+        processed_inputs = processor(images=raw_input, return_tensors="pt")
+        outputs = onnx_model(**processed_inputs)
+
+        # Check device and output format
+        assert onnx_model.device.type == "cpu"
+        assert isinstance(outputs.last_hidden_state, torch.Tensor)
+        features = outputs.last_hidden_state.detach().cpu().numpy().tolist()
+        assert all(isinstance(item, float) for row in features for inner in row for item in inner)
+        gc.collect()
+
+    @parameterized.expand(
+        grid_parameters(
+            {"model_arch": SUPPORTED_ARCHITECTURES, "provider": ["CUDAExecutionProvider", "TensorrtExecutionProvider"]}
+        )
+    )
+    @require_torch_gpu
+    @pytest.mark.cuda_ep_test
+    @pytest.mark.trt_ep_test
+    def test_inference_on_gpu(self, test_name: str, model_arch: str, provider: str):
+        if provider == "TensorrtExecutionProvider" and model_arch != self.__class__.SUPPORTED_ARCHITECTURES[0]:
+            self.skipTest("testing a single arch for TensorrtExecutionProvider")
+
+        model_args = {"test_name": model_arch, "model_arch": model_arch}
+        self._setup(model_args)
+
+        model_id = MODEL_NAMES[model_arch]
+        onnx_model = ORTModelForFeatureExtraction.from_pretrained(self.onnx_model_dirs[model_arch], provider=provider)
+        processor = get_preprocessor(model_id)
+        raw_input = self.get_raw_input(model_arch)
+        processed_inputs = processor(images=raw_input, return_tensors="pt").to("cuda")
+        outputs = onnx_model(**processed_inputs)
+
+        # Check device and output format
+        assert onnx_model.device.type == "cuda"
+        assert isinstance(outputs.last_hidden_state, torch.Tensor)
+        features = outputs.last_hidden_state.detach().cpu().numpy().tolist()
+        assert all(isinstance(item, float) for row in features for inner in row for item in inner)
+
+        gc.collect()
+
+    @parameterized.expand(SUPPORTED_ARCHITECTURES)
+    @require_torch_gpu
+    @pytest.mark.cuda_ep_test
+    def test_compare_to_io_binding(self, model_arch):
+        model_args = {"test_name": model_arch, "model_arch": model_arch}
+        self._setup(model_args)
+
+        model_id = MODEL_NAMES[model_arch]
+        onnx_model = ORTModelForFeatureExtraction.from_pretrained(
+            self.onnx_model_dirs[model_arch], use_io_binding=False, provider="CUDAExecutionProvider"
+        )
+        io_model = ORTModelForFeatureExtraction.from_pretrained(
+            self.onnx_model_dirs[model_arch], use_io_binding=True, provider="CUDAExecutionProvider"
+        )
+
+        self.assertFalse(onnx_model.use_io_binding)
+        self.assertTrue(io_model.use_io_binding)
+
+        processor = get_preprocessor(model_id)
+        raw_input = self.get_raw_input(model_arch)
+        tokens = processor(images=[raw_input, raw_input], return_tensors="pt").to("cuda")
+
+        onnx_outputs = onnx_model(**tokens)
+        io_outputs = io_model(**tokens)
+
+        self.assertTrue("last_hidden_state" in io_outputs)
+        self.assertIsInstance(io_outputs.last_hidden_state, torch.Tensor)
+
+        # compare tensor outputs
+        torch.testing.assert_close(
+            onnx_outputs.last_hidden_state, io_outputs.last_hidden_state, atol=self.ATOL, rtol=self.RTOL
+        )
+
+        gc.collect()
+
+
 class ORTModelForMultipleChoiceIntegrationTest(ORTModelTestMixin):
     # Multiple Choice tests are conducted on different models due to mismatch size in model's classifier
     SUPPORTED_ARCHITECTURES = [


Reviews:
user: login: IlyasMoutawwakil
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-09T13:14:49Z
user: login: IlyasMoutawwakil
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-09T13:15:42Z
user: login: IlyasMoutawwakil
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-09T14:10:10Z
user: login: IlyasMoutawwakil
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-09T14:11:21Z
user: login: Abdennacer-Badaoui
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-09T14:13:59Z

Review comments:
diff_hunk: @@ -1061,18 +1061,34 @@ def forward(
         token_type_ids: Optional[Union[torch.Tensor, np.ndarray]] = None,
         **kwargs,
     ):
-        use_torch = isinstance(input_ids, torch.Tensor)
+        # Determine the tensor type from any available tensor input
+        tensor_inputs = [input_ids, attention_mask, token_type_ids] + [
+            v for k, v in kwargs.items() if isinstance(v, (torch.Tensor, np.ndarray))
+        ]
+        first_tensor = next((t for t in tensor_inputs if t is not None), None)
+        use_torch = isinstance(first_tensor, torch.Tensor) if first_tensor is not None else False
         self.raise_on_numpy_input_io_binding(use_torch)
 
         if token_type_ids is None and "token_type_ids" in self.input_names:
             token_type_ids = torch.zeros_like(input_ids) if use_torch else np.zeros_like(input_ids)
 
+        # Build model_inputs dictionary
         model_inputs = {
-            "input_ids": input_ids,
-            "attention_mask": attention_mask,
-            "token_type_ids": token_type_ids,
+            k: v
+            for k, v in {
+                "input_ids": input_ids,
+                "attention_mask": attention_mask,
+                "token_type_ids": token_type_ids,
+                **kwargs,
+            }.items()
+            if k in self.input_names and v is not None
         }
 
+        # Validate that we have all required inputs
+        for input_name in self.input_names:
+            if input_name not in model_inputs:
+                raise ValueError(f"Input {input_name} is required by model but not provided.")
path: optimum/onnxruntime/modeling_ort.py
user: login: IlyasMoutawwakil
body: this part is handled by _prepare_onnx_inputs and _prepare_io_binding
created_at: 2025-05-09T13:14:49Z
author_association: MEMBER
diff_hunk: @@ -1061,18 +1061,34 @@ def forward(
         token_type_ids: Optional[Union[torch.Tensor, np.ndarray]] = None,
         **kwargs,
     ):
-        use_torch = isinstance(input_ids, torch.Tensor)
+        # Determine the tensor type from any available tensor input
+        tensor_inputs = [input_ids, attention_mask, token_type_ids] + [
+            v for k, v in kwargs.items() if isinstance(v, (torch.Tensor, np.ndarray))
+        ]
+        first_tensor = next((t for t in tensor_inputs if t is not None), None)
+        use_torch = isinstance(first_tensor, torch.Tensor) if first_tensor is not None else False
         self.raise_on_numpy_input_io_binding(use_torch)
 
         if token_type_ids is None and "token_type_ids" in self.input_names:
             token_type_ids = torch.zeros_like(input_ids) if use_torch else np.zeros_like(input_ids)
 
+        # Build model_inputs dictionary
         model_inputs = {
-            "input_ids": input_ids,
-            "attention_mask": attention_mask,
-            "token_type_ids": token_type_ids,
+            k: v
+            for k, v in {
+                "input_ids": input_ids,
+                "attention_mask": attention_mask,
+                "token_type_ids": token_type_ids,
+                **kwargs,
+            }.items()
+            if k in self.input_names and v is not None
path: optimum/onnxruntime/modeling_ort.py
user: login: IlyasMoutawwakil
body: no need to filter them here since _prepare_onnx_inputs and _prepare_io_binding will raise error on missing inputs
created_at: 2025-05-09T13:15:42Z
author_association: MEMBER
diff_hunk: @@ -1059,18 +1059,37 @@ def forward(
         input_ids: Optional[Union[torch.Tensor, np.ndarray]] = None,
         attention_mask: Optional[Union[torch.Tensor, np.ndarray]] = None,
         token_type_ids: Optional[Union[torch.Tensor, np.ndarray]] = None,
-        **kwargs,
+        position_ids: Optional[Union[torch.Tensor, np.ndarray]] = None,
+        pixel_values: Optional[Union[torch.Tensor, np.ndarray]] = None,
+        input_features: Optional[Union[torch.Tensor, np.ndarray]] = None,
+        input_values: Optional[Union[torch.Tensor, np.ndarray]] = None,
     ):
-        use_torch = isinstance(input_ids, torch.Tensor)
+        # Determine the tensor type from any available tensor input
+        tensor_inputs = [
+            input_ids,
+            attention_mask,
+            token_type_ids,
+            position_ids,
+            pixel_values,
+            input_features,
+            input_values,
+        ]
+        first_tensor = next((t for t in tensor_inputs if t is not None), None)
+        use_torch = isinstance(first_tensor, torch.Tensor) if first_tensor is not None else False
path: optimum/onnxruntime/modeling_ort.py
user: login: IlyasMoutawwakil
body: no need to check for first_tensor is None

```suggestion
        first_tensor = next(filter(None, tensor_inputs))
        use_torch = isinstance(first_tensor, torch.Tensor)
```
created_at: 2025-05-09T14:10:10Z
author_association: MEMBER
diff_hunk: @@ -2160,14 +2177,26 @@ def test_pipeline_ort_model(self, model_arch):
 
         model_id = MODEL_NAMES[model_arch]
         onnx_model = ORTModelForFeatureExtraction.from_pretrained(self.onnx_model_dirs[model_arch])
-        tokenizer = get_preprocessor(model_id)
-        pipe = pipeline("feature-extraction", model=onnx_model, tokenizer=tokenizer)
-        text = "My Name is Philipp and i live in Germany."
-        outputs = pipe(text)
-
-        # compare model output class
-        self.assertEqual(pipe.device, onnx_model.device)
-        self.assertTrue(all(all(isinstance(item, float) for item in row) for row in outputs[0]))
+        processor = get_preprocessor(model_id)
+        if self.is_image_model(model_arch):
+            # For image models, bypass the pipeline and test directly because feature extraction (in optimum/pipelines/pipelines_base.py) is only supported for text at the moment
+            # Change this once feature extraction pipelines is supported for image models
+            raw_input = self.get_raw_input(model_arch)
+            processed_inputs = processor(images=raw_input, return_tensors="pt")
+            outputs = onnx_model(**processed_inputs)
+
+            # Check device and output format
+            assert onnx_model.device.type == "cpu"
+            assert isinstance(outputs.last_hidden_state, torch.Tensor)
+            features = outputs.last_hidden_state.detach().cpu().numpy().tolist()
+            assert all(isinstance(item, float) for row in features for inner in row for item in inner)
+        else:
+            # For text models, use the pipeline and check model device and compare model output class
+            pipe = pipeline("feature-extraction", model=onnx_model, tokenizer=processor)
+            raw_input = self.get_raw_input(model_arch)
+            outputs = pipe(raw_input)
+            assert pipe.device == onnx_model.device
+            assert all(all(isinstance(item, float) for item in row) for row in outputs[0])
path: tests/onnxruntime/test_modeling.py
user: login: IlyasMoutawwakil
body: the test `test_pipeline_ort_model` needs to check using pipeline
created_at: 2025-05-09T14:11:21Z
author_association: MEMBER
diff_hunk: @@ -2160,14 +2177,26 @@ def test_pipeline_ort_model(self, model_arch):
 
         model_id = MODEL_NAMES[model_arch]
         onnx_model = ORTModelForFeatureExtraction.from_pretrained(self.onnx_model_dirs[model_arch])
-        tokenizer = get_preprocessor(model_id)
-        pipe = pipeline("feature-extraction", model=onnx_model, tokenizer=tokenizer)
-        text = "My Name is Philipp and i live in Germany."
-        outputs = pipe(text)
-
-        # compare model output class
-        self.assertEqual(pipe.device, onnx_model.device)
-        self.assertTrue(all(all(isinstance(item, float) for item in row) for row in outputs[0]))
+        processor = get_preprocessor(model_id)
+        if self.is_image_model(model_arch):
+            # For image models, bypass the pipeline and test directly because feature extraction (in optimum/pipelines/pipelines_base.py) is only supported for text at the moment
+            # Change this once feature extraction pipelines is supported for image models
+            raw_input = self.get_raw_input(model_arch)
+            processed_inputs = processor(images=raw_input, return_tensors="pt")
+            outputs = onnx_model(**processed_inputs)
+
+            # Check device and output format
+            assert onnx_model.device.type == "cpu"
+            assert isinstance(outputs.last_hidden_state, torch.Tensor)
+            features = outputs.last_hidden_state.detach().cpu().numpy().tolist()
+            assert all(isinstance(item, float) for row in features for inner in row for item in inner)
+        else:
+            # For text models, use the pipeline and check model device and compare model output class
+            pipe = pipeline("feature-extraction", model=onnx_model, tokenizer=processor)
+            raw_input = self.get_raw_input(model_arch)
+            outputs = pipe(raw_input)
+            assert pipe.device == onnx_model.device
+            assert all(all(isinstance(item, float) for item in row) for row in outputs[0])
path: tests/onnxruntime/test_modeling.py
user: login: Abdennacer-Badaoui
body: Yes, since the feature-extraction pipelines currently support only text, should I work on extending them?
created_at: 2025-05-09T14:13:59Z
author_association: MEMBER

PR comments:
user: login: HuggingFaceDocBuilderDev
created_at: 2025-05-09T10:55:00Z
body: The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/optimum/pr_2255). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.
author_association: NONE
user: login: IlyasMoutawwakil
created_at: 2025-05-09T14:13:15Z
body: No need to change the testing suite for this, as it's gonna be undergoing some refactoring
author_association: MEMBER