Repository: huggingface/optimum

Issue:
Add IOBinding for ORTModelForVision2Seq
### Feature request

Add IOBinding support for ORTModelForVision2Seq

### Motivation

Priority model for inference

### Your contribution




Discussion:
Also tagging @IlyasMoutawwakil, as he tested with it yesterday!
@mht-sharma wile trying to solve it I found that the problem was in an output axis name that couldn't be inferred and wasn't in `dimensions`. It's name is `encoder_sequence_length`. As a quick patch I added it to dimensions when possible [here](https://github.com/huggingface/optimum/blob/e32ba52e4025692e1979cbbf0b4a55c02a3f2ac6/optimum/onnxruntime/modeling_ort.py#L769):
```python
if ("height" in dimensions) and ("width" in dimensions) and hasattr(self.config, "encoder"):
    dimensions["encoder_sequence_length"] = dimensions["height"] * dimensions["width"] // self.config.encoder.hidden_size
```
Not sure this is the best way to do it.

Resolving patch:
diff --git a/optimum/exporters/tasks.py b/optimum/exporters/tasks.py
--- a/optimum/exporters/tasks.py
+++ b/optimum/exporters/tasks.py
@@ -481,6 +481,13 @@ class TasksManager:
             onnx="DistilBertOnnxConfig",
             tflite="DistilBertTFLiteConfig",
         ),
+        "donut": supported_tasks_mapping(
+            "image-to-text",
+            "image-to-text-with-past",
+            "document-question-answering",
+            "document-question-answering-with-past",
+            onnx="VisionEncoderDecoderOnnxConfig",
+        ),
         "donut-swin": supported_tasks_mapping(
             "feature-extraction",
             onnx="DonutSwinOnnxConfig",
diff --git a/optimum/onnxruntime/modeling_seq2seq.py b/optimum/onnxruntime/modeling_seq2seq.py
--- a/optimum/onnxruntime/modeling_seq2seq.py
+++ b/optimum/onnxruntime/modeling_seq2seq.py
@@ -417,8 +417,13 @@ def forward(
         self.parent_model.raise_on_numpy_input_io_binding(use_torch)
 
         if self.parent_model.device.type == "cuda" and self.parent_model.use_io_binding:
+            known_output_shapes = self.compute_encoder_known_output_shapes(pixel_values)
+
             io_binding, output_shapes, output_buffers = self.parent_model._prepare_io_binding(
-                self.session, pixel_values, ordered_input_names=self._ordered_input_names
+                self.session,
+                pixel_values,
+                known_output_shapes=known_output_shapes,
+                ordered_input_names=self._ordered_input_names,
             )
 
             io_binding.synchronize_inputs()
@@ -440,6 +445,34 @@ def forward(
 
         return BaseModelOutput(last_hidden_state=last_hidden_state)
 
+    def compute_encoder_known_output_shapes(self, pixel_values: torch.FloatTensor) -> Dict[str, List[int]]:
+        if self.normalized_config.model_type == "vit":
+            # for vit models
+            encoder_sequence_length = (
+                self.normalized_config.image_size // self.normalized_config.config.patch_size
+            ) ** 2 + 1  # plus cls token
+        elif self.normalized_config.config.model_type == "donut-swin":
+            # for donut-swin models
+            encoder_sequence_length = (
+                self.normalized_config.config.image_size[0]
+                * self.normalized_config.config.image_size[1]
+                // self.normalized_config.config.hidden_size
+            )
+        else:
+            raise ValueError(
+                f"Unsupported encoder model type {self.normalized_config.config.model_type} for ORTForVisionSeq2Seq with IOBinding."
+                "Currently supported models are vit and donut-swin."
+                "Please submit a PR to add support for this model type."
+            )
+
+        return {
+            "last_hidden_state": [
+                pixel_values.shape[0],  # batch_size
+                encoder_sequence_length,  # encoder_sequence_length
+                self.normalized_config.config.hidden_size,  # hidden_size
+            ]
+        }
+
 
 class ORTEncoderForPix2Struct(ORTEncoder):
     """


Test patch:
diff --git a/tests/exporters/exporters_utils.py b/tests/exporters/exporters_utils.py
--- a/tests/exporters/exporters_utils.py
+++ b/tests/exporters/exporters_utils.py
@@ -51,6 +51,7 @@
     "deberta": "hf-internal-testing/tiny-random-DebertaModel",
     "deberta-v2": "hf-internal-testing/tiny-random-DebertaV2Model",
     "deit": "hf-internal-testing/tiny-random-DeiTModel",
+    "donut": "fxmarty/tiny-doc-qa-vision-encoder-decoder",
     "donut-swin": "hf-internal-testing/tiny-random-DonutSwinModel",
     "detr": "hf-internal-testing/tiny-random-DetrModel",  # hf-internal-testing/tiny-random-detr is larger
     "distilbert": "hf-internal-testing/tiny-random-DistilBertModel",
diff --git a/tests/onnxruntime/test_modeling.py b/tests/onnxruntime/test_modeling.py
--- a/tests/onnxruntime/test_modeling.py
+++ b/tests/onnxruntime/test_modeling.py
@@ -3957,8 +3957,7 @@ def test_compare_generation_to_io_binding(
 
 
 class ORTModelForVision2SeqIntegrationTest(ORTModelTestMixin):
-    # TODO: speech_to_text should be tested
-    SUPPORTED_ARCHITECTURES = ["vision-encoder-decoder", "trocr"]
+    SUPPORTED_ARCHITECTURES = ["vision-encoder-decoder", "trocr", "donut"]
 
     FULL_GRID = {
         "model_arch": SUPPORTED_ARCHITECTURES,
@@ -4157,7 +4156,6 @@ def test_pipeline_on_gpu(self, test_name: str, model_arch: str, use_cache: bool)
             "test_name": test_name,
             "model_arch": model_arch,
             "use_cache": use_cache,
-            "use_io_binding": False,
         }
         self._setup(model_args)
 
@@ -4225,6 +4223,93 @@ def test_compare_with_and_without_past_key_values(self, model_arch: str):
                 f" speedup: {without_pkv_timer.elapsed / with_pkv_timer.elapsed:.3f}",
             )
 
+    @parameterized.expand(
+        grid_parameters({"model_arch": SUPPORTED_ARCHITECTURES, "use_cache": [True], "use_merged": [False, True]})
+    )
+    @require_torch_gpu
+    @pytest.mark.gpu_test
+    def test_compare_to_io_binding(self, test_name: str, model_arch: str, use_cache: bool, use_merged: bool):
+        if use_cache is False and use_merged is True:
+            self.skipTest("use_cache=False, use_merged=True are uncompatible")
+
+        model_args = {
+            "test_name": test_name,
+            "model_arch": model_arch,
+            "use_cache": use_cache,
+            "use_merged": use_merged,
+        }
+        self._setup(model_args)
+
+        model_id = MODEL_NAMES[model_arch]
+        onnx_model = ORTModelForVision2Seq.from_pretrained(self.onnx_model_dirs[test_name], use_io_binding=False).to(
+            "cuda"
+        )
+        io_model = ORTModelForVision2Seq.from_pretrained(self.onnx_model_dirs[test_name], use_io_binding=True).to(
+            "cuda"
+        )
+
+        self.assertFalse(onnx_model.use_io_binding)
+        self.assertTrue(io_model.use_io_binding)
+
+        feature_extractor, tokenizer = self._get_preprocessors(model_id)
+
+        data = self._get_sample_image()
+        pixel_values = feature_extractor([data] * 2, return_tensors="pt").pixel_values.to("cuda")
+
+        decoder_start_token_id = onnx_model.config.decoder.bos_token_id
+        decoder_input_ids = torch.full((2, 1), decoder_start_token_id, dtype=torch.long).to("cuda")
+
+        onnx_outputs = onnx_model(pixel_values=pixel_values, decoder_input_ids=decoder_input_ids)
+        io_outputs = io_model(pixel_values=pixel_values, decoder_input_ids=decoder_input_ids)
+
+        self.assertTrue("logits" in io_outputs)
+        self.assertIsInstance(io_outputs.logits, torch.Tensor)
+
+        # compare tensor outputs
+        self.assertTrue(torch.equal(onnx_outputs.logits, io_outputs.logits))
+
+        gc.collect()
+
+    @parameterized.expand(
+        grid_parameters({"model_arch": SUPPORTED_ARCHITECTURES, "use_cache": [True], "use_merged": [False, True]})
+    )
+    @require_torch_gpu
+    @pytest.mark.gpu_test
+    def test_compare_generation_to_io_binding(
+        self, test_name: str, model_arch: str, use_cache: bool, use_merged: bool
+    ):
+        if use_cache is False and use_merged is True:
+            self.skipTest("use_cache=False, use_merged=True are uncompatible")
+
+        model_args = {
+            "test_name": test_name,
+            "model_arch": model_arch,
+            "use_cache": use_cache,
+            "use_merged": use_merged,
+        }
+        self._setup(model_args)
+
+        model_id = MODEL_NAMES[model_arch]
+        onnx_model = ORTModelForVision2Seq.from_pretrained(self.onnx_model_dirs[test_name], use_io_binding=False).to(
+            "cuda"
+        )
+        io_model = ORTModelForVision2Seq.from_pretrained(self.onnx_model_dirs[test_name], use_io_binding=True).to(
+            "cuda"
+        )
+
+        feature_extractor, tokenizer = self._get_preprocessors(model_id)
+
+        data = self._get_sample_image()
+        features = feature_extractor(data, return_tensors="pt").to("cuda")
+
+        onnx_outputs = onnx_model.generate(**features, num_beams=5)
+        io_outputs = io_model.generate(**features, num_beams=5)
+
+        # compare tensor outputs
+        self.assertTrue(torch.equal(onnx_outputs, io_outputs))
+
+        gc.collect()
+
 
 class ORTModelForCustomTasksIntegrationTest(ORTModelTestMixin):
     SUPPORTED_ARCHITECTURES_WITH_MODEL_ID = {
diff --git a/tests/onnxruntime/utils_onnxruntime_tests.py b/tests/onnxruntime/utils_onnxruntime_tests.py
--- a/tests/onnxruntime/utils_onnxruntime_tests.py
+++ b/tests/onnxruntime/utils_onnxruntime_tests.py
@@ -47,6 +47,7 @@
     "deberta": "hf-internal-testing/tiny-random-DebertaModel",
     "deberta_v2": "hf-internal-testing/tiny-random-DebertaV2Model",
     "deit": "hf-internal-testing/tiny-random-DeiTModel",
+    "donut": "fxmarty/tiny-doc-qa-vision-encoder-decoder",
     "detr": "hf-internal-testing/tiny-random-detr",
     "distilbert": "hf-internal-testing/tiny-random-DistilBertModel",
     "electra": "hf-internal-testing/tiny-random-ElectraModel",