Repository: huggingface/optimum

Issue:
Cannot convert owlvit-base-patch32 model to ONNX and run inference
### System Info

```shell
Optimum version: 1.9.1
Python version: 3.11.3
OS: MacOS
```


### Who can help?

@mich

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [X] My own task or dataset (give details below)

### Reproduction

When using the CLI command
`optimum-cli export onnx --model google/owlvit-base-patch32 --task zero-shot-object-detection object_detection/owlvit_onnx` 
I'm able to get a converted ONNX format. Then, when using the following code to perform inference with the converted model:
`checkpoint = "google/owlvit-base-patch32"`
`processor = AutoProcessor.from_pretrained(checkpoint)`

`image = skimage.data.astronaut()`
`image = Image.fromarray(np.uint8(image)).convert("RGB")`
`text_queries = ["human face", "rocket", "nasa badge", "star-spangled banner", "woman", "smile", "hair", 'human head', 'human eye']`

`np_inputs = processor(text=text_queries, images=image, return_tensors="np")`
`session = ort.InferenceSession("object_detection/owlvit_onnx/model.onnx")`

`out =session.run(['logits', 'pred_boxes', 'text_embeds', 'image_embeds'], np_inputs)`

I get the following error: 
`RuntimeException: [ONNXRuntimeError] : 6 : RUNTIME_EXCEPTION : Non-zero status code returned while running Reshape node. Name:'/Reshape_3' Status Message: /Users/runner/work/1/s/onnxruntime/core/providers/cpu/tensor/reshape_helper.h:41 onnxruntime::ReshapeHelper::ReshapeHelper(const onnxruntime::TensorShape &, onnxruntime::TensorShapeVector &, bool) gsl::narrow_cast(input_shape.Size()) == size was false. The input tensor cannot be reshaped to the requested shape. Input shape:{9,16}, requested shape:{2,4,16}`

Now it seems to be related to some input being wrong, but I cannot get what is wrong. The pre-processing step is the same as for the HF model, the only difference being instead of returning "pt" tensors I'm returning "np" so it can work with ONNX. Here are my input shapes:

input_ids: (9, 16)
attention_mask: (9, 16)
pixel_values: (1, 3, 768, 768)

Thanks in advance!

### Expected behavior

Inference to run successfully and outputs to be very similar to that of the original torch model. 


Discussion:

Resolving patch:
diff --git a/optimum/exporters/onnx/model_configs.py b/optimum/exporters/onnx/model_configs.py
--- a/optimum/exporters/onnx/model_configs.py
+++ b/optimum/exporters/onnx/model_configs.py
@@ -841,6 +841,15 @@ class OwlViTOnnxConfig(CLIPOnnxConfig):
     ATOL_FOR_VALIDATION = 1e-4
     MIN_TORCH_VERSION = version.parse("2.1")
 
+    def __init__(self, config: "PretrainedConfig", task: str = "feature-extraction"):
+        super().__init__(config, task)
+        if task == "zero-shot-object-detection":
+            logger.warning(
+                "The batch size of this model will not be dynamic because non-maximum suppression is performed. "
+                "Make sure to export the model with the same batch size as the one you will use at inference "
+                "with `--batch_size N`."
+            )
+
     @property
     def outputs(self) -> Dict[str, Dict[int, str]]:
         outputs = {}
@@ -848,10 +857,10 @@ def outputs(self) -> Dict[str, Dict[int, str]]:
             outputs["logits_per_image"] = {0: "image_batch_size", 1: "text_batch_size"}
             outputs["logits_per_text"] = {0: "text_batch_size", 1: "image_batch_size"}
         elif self.task == "zero-shot-object-detection":
-            outputs["logits"] = {0: "batch_size", 1: "num_queries"}
-            outputs["pred_boxes"] = {0: "batch_size", 1: "num_queries"}
+            outputs["logits"] = {0: "image_batch_size", 2: "num_queries"}
+            outputs["pred_boxes"] = {0: "image_batch_size", 1: "num_boxes"}
 
-        outputs["text_embeds"] = {0: "text_batch_size"}
+        outputs["text_embeds"] = {0: "text_batch_size", 1: "max_text_queries"}
         outputs["image_embeds"] = {0: "image_batch_size"}
         return outputs
 


Test patch: