Add ONNX export for owlv2

Author: xenova

optimum/exporters/onnx/model_configs.py

@@ -1085,6 +1085,10 @@ def outputs(self) -> Dict[str, Dict[int, str]]:
         return outputs
 
 
+class OwlV2OnnxConfig(OwlViTOnnxConfig):
+    DEFAULT_ONNX_OPSET = 14
+
+
 class LayoutLMOnnxConfig(TextAndVisionOnnxConfig):
     NORMALIZED_CONFIG_CLASS = NormalizedTextConfig.with_args(
         allow_new=True,

optimum/exporters/tasks.py

@@ -808,6 +808,11 @@ class TasksManager:
             "token-classification",
             onnx="NystromformerOnnxConfig",
         ),
+        "owlv2": supported_tasks_mapping(
+            "feature-extraction",
+            "zero-shot-object-detection",
+            onnx="OwlV2OnnxConfig",
+        ),
         "owlvit": supported_tasks_mapping(
             "feature-extraction",
             "zero-shot-object-detection",

optimum/utils/normalized_config.py

@@ -199,6 +199,7 @@ class NormalizedConfigManager:
         'levit',
         'mobilebert',
         'mobilevit',
+        'owlv2',
         'owlvit',
         'perceiver',
         'roformer',

tests/exporters/exporters_utils.py

@@ -121,6 +121,7 @@
     "mt5": "lewtun/tiny-random-mt5",
     "nystromformer": "hf-internal-testing/tiny-random-NystromformerModel",
     "opt": "hf-internal-testing/tiny-random-OPTModel",
+    "owlv2": "hf-internal-testing/tiny-random-Owlv2Model",
     "owlvit": "hf-tiny-model-private/tiny-random-OwlViTModel",
     "pegasus": "hf-internal-testing/tiny-random-PegasusModel",
     "perceiver": {
@@ -241,6 +242,7 @@
     "mpt": "mosaicml/mpt-7b",
     "mt5": "lewtun/tiny-random-mt5",  # Not using google/mt5-small because it takes too much time for testing.
     "nystromformer": "hf-internal-testing/tiny-random-NystromformerModel",
+    "owlv2": "google/owlv2-base-patch16",
     "owlvit": "google/owlvit-base-patch32",
     "perceiver": "hf-internal-testing/tiny-random-PerceiverModel",  # Not using deepmind/language-perceiver because it takes too much time for testing.
     # "rembert": "google/rembert",