ViTPose export fix (#2192)

* Remove feature extraction support for vitpose * remove unnecessary
huggingface · Feb 19, 2025 · 3d8957c · 3d8957c
1 parent 0245619
commit 3d8957c
Show file tree

Hide file tree

Showing 2 changed files with 1 addition and 13 deletions.
diff --git a/optimum/exporters/tasks.py b/optimum/exporters/tasks.py
@@ -1242,7 +1242,7 @@ class TasksManager:
             "image-classification",
             onnx="VitMSNOnnxConfig",
         ),
-        "vitpose": supported_tasks_mapping("feature-extraction", "keypoint-detection", onnx="VitPoseOnnxConfig"),
+        "vitpose": supported_tasks_mapping("keypoint-detection", onnx="VitPoseOnnxConfig"),
         "vits": supported_tasks_mapping(
             "text-to-audio",
             onnx="VitsOnnxConfig",

diff --git a/optimum/utils/input_generators.py b/optimum/utils/input_generators.py
@@ -1622,12 +1622,6 @@ def __init__(
             self.height = preprocessor.crop_size.get("height", self.height)
             self.width = preprocessor.crop_size.get("width", self.width)
 
-    def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"):
-        input_ = super().generate(
-            input_name=input_name, framework=framework, int_dtype=int_dtype, float_dtype=float_dtype
-        )
-        return input_
-
 
 class DummyVisionStaticInputGenerator(DummyVisionInputGenerator):
     def __init__(
@@ -1657,12 +1651,6 @@ def __init__(
             self.height = preprocessor.size.get("height", self.height)
             self.width = preprocessor.size.get("width", self.width)
 
-    def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"):
-        input_ = super().generate(
-            input_name=input_name, framework=framework, int_dtype=int_dtype, float_dtype=float_dtype
-        )
-        return input_
-
 
 class PerceiverDummyInputGenerator(DummyVisionStaticInputGenerator):
     pass