diff --git a/optimum/exporters/onnx/model_configs.py b/optimum/exporters/onnx/model_configs.py index 9bbfacf235..e6518d5cdd 100644 --- a/optimum/exporters/onnx/model_configs.py +++ b/optimum/exporters/onnx/model_configs.py @@ -1092,6 +1092,10 @@ def outputs(self) -> Dict[str, Dict[int, str]]: return outputs +class OwlV2OnnxConfig(OwlViTOnnxConfig): + pass + + class LayoutLMOnnxConfig(TextAndVisionOnnxConfig): NORMALIZED_CONFIG_CLASS = NormalizedTextConfig.with_args( allow_new=True, diff --git a/optimum/exporters/tasks.py b/optimum/exporters/tasks.py index 57ee54c2c7..f12af9bf1d 100644 --- a/optimum/exporters/tasks.py +++ b/optimum/exporters/tasks.py @@ -832,6 +832,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", diff --git a/optimum/utils/normalized_config.py b/optimum/utils/normalized_config.py index 2f705aed2c..0a1d19fa75 100644 --- a/optimum/utils/normalized_config.py +++ b/optimum/utils/normalized_config.py @@ -199,6 +199,7 @@ class NormalizedConfigManager: 'levit', 'mobilebert', 'mobilevit', + 'owlv2', 'owlvit', 'perceiver', 'roformer', diff --git a/tests/exporters/exporters_utils.py b/tests/exporters/exporters_utils.py index 948bb7001e..79b60641ad 100644 --- a/tests/exporters/exporters_utils.py +++ b/tests/exporters/exporters_utils.py @@ -122,6 +122,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": { @@ -243,6 +244,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",