Skip to content

Commit 3d8957c

Browse files
authored
ViTPose export fix (#2192)
* Remove feature extraction support for vitpose * remove unnecessary
1 parent 0245619 commit 3d8957c

File tree

2 files changed

+1
-13
lines changed

2 files changed

+1
-13
lines changed

optimum/exporters/tasks.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -1242,7 +1242,7 @@ class TasksManager:
12421242
"image-classification",
12431243
onnx="VitMSNOnnxConfig",
12441244
),
1245-
"vitpose": supported_tasks_mapping("feature-extraction", "keypoint-detection", onnx="VitPoseOnnxConfig"),
1245+
"vitpose": supported_tasks_mapping("keypoint-detection", onnx="VitPoseOnnxConfig"),
12461246
"vits": supported_tasks_mapping(
12471247
"text-to-audio",
12481248
onnx="VitsOnnxConfig",

optimum/utils/input_generators.py

-12
Original file line numberDiff line numberDiff line change
@@ -1622,12 +1622,6 @@ def __init__(
16221622
self.height = preprocessor.crop_size.get("height", self.height)
16231623
self.width = preprocessor.crop_size.get("width", self.width)
16241624

1625-
def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"):
1626-
input_ = super().generate(
1627-
input_name=input_name, framework=framework, int_dtype=int_dtype, float_dtype=float_dtype
1628-
)
1629-
return input_
1630-
16311625

16321626
class DummyVisionStaticInputGenerator(DummyVisionInputGenerator):
16331627
def __init__(
@@ -1657,12 +1651,6 @@ def __init__(
16571651
self.height = preprocessor.size.get("height", self.height)
16581652
self.width = preprocessor.size.get("width", self.width)
16591653

1660-
def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"):
1661-
input_ = super().generate(
1662-
input_name=input_name, framework=framework, int_dtype=int_dtype, float_dtype=float_dtype
1663-
)
1664-
return input_
1665-
16661654

16671655
class PerceiverDummyInputGenerator(DummyVisionStaticInputGenerator):
16681656
pass

0 commit comments

Comments
 (0)