Repository: huggingface/optimum

Issue:
ValueError: Unrecognized configuration class <class 'transformers.models.mobilevit.configuration_mobilevit.MobileViTConfig'> for this kind of AutoModel: AutoModelForImageSegmentation.
### System Info

```shell
- `optimum` version: 1.8.8.dev0
- `transformers` version: 4.30.2
- Platform: Linux-5.15.107+-x86_64-with-glibc2.31
- Python version: 3.10.12
- Huggingface_hub version: 0.15.1
- PyTorch version (GPU?): 2.0.1+cu118 (cuda availabe: False)
- Tensorflow version (GPU?): not installed (cuda availabe: NA)
```


### Who can help?

@michaelbenayoun 

### 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

Running `optimum-cli export onnx -m apple/deeplabv3-mobilevit-small out`

results in:
```
Framework not specified. Using pt to export to ONNX.
Traceback (most recent call last):
  File "/usr/local/bin/optimum-cli", line 8, in <module>
    sys.exit(main())
  File "/usr/local/lib/python3.10/dist-packages/optimum/commands/optimum_cli.py", line 163, in main
    service.run()
  File "/usr/local/lib/python3.10/dist-packages/optimum/commands/export/onnx.py", line 219, in run
    main_export(
  File "/usr/local/lib/python3.10/dist-packages/optimum/exporters/onnx/__main__.py", line 184, in main_export
    model = TasksManager.get_model_from_task(
  File "/usr/local/lib/python3.10/dist-packages/optimum/exporters/tasks.py", line 1470, in get_model_from_task
    model = model_class.from_pretrained(model_name_or_path, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/transformers/models/auto/auto_factory.py", line 487, in from_pretrained
    raise ValueError(
ValueError: Unrecognized configuration class <class 'transformers.models.mobilevit.configuration_mobilevit.MobileViTConfig'> for this kind of AutoModel: AutoModelForImageSegmentation.
Model type should be one of DetrConfig.
```

### Expected behavior

The model should export correctly.


Discussion:

Resolving patch:
diff --git a/optimum/exporters/onnx/base.py b/optimum/exporters/onnx/base.py
--- a/optimum/exporters/onnx/base.py
+++ b/optimum/exporters/onnx/base.py
@@ -133,14 +133,7 @@ class OnnxConfig(ExportConfig, ABC):
         "feature-extraction": OrderedDict({"last_hidden_state": {0: "batch_size", 1: "sequence_length"}}),
         "fill-mask": OrderedDict({"logits": {0: "batch_size", 1: "sequence_length"}}),
         "image-classification": OrderedDict({"logits": {0: "batch_size"}}),
-        # TODO: Is this the same thing as semantic-segmentation?
-        "image-segmentation": OrderedDict(
-            {
-                "logits": {0: "batch_size", 1: "num_queries"},
-                "pred_boxes": {0: "batch_size", 1: "num_queries"},
-                "pred_masks": {0: "batch_size", 1: "num_queries"},
-            }
-        ),
+        "image-segmentation": OrderedDict({"logits": {0: "batch_size", 1: "num_labels", 2: "height", 3: "width"}}),
         "image-to-text": OrderedDict({"logits": {0: "batch_size", 1: "sequence_length"}}),
         "mask-generation": OrderedDict({"logits": {0: "batch_size"}}),
         "masked-im": OrderedDict(
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
@@ -572,7 +572,7 @@ class ConvNextOnnxConfig(ViTOnnxConfig):
 
 
 class MobileViTOnnxConfig(ViTOnnxConfig):
-    pass
+    ATOL_FOR_VALIDATION = 1e-4
 
 
 class RegNetOnnxConfig(ViTOnnxConfig):
@@ -588,9 +588,14 @@ class DetrOnnxConfig(ViTOnnxConfig):
     DEFAULT_ONNX_OPSET = 12
 
     @property
-    def inputs(self) -> Dict[str, Dict[int, str]]:
-        # TODO: is pixel mask needed?
-        return {**super().inputs, "pixel_mask": {0: "batch_size"}}
+    def outputs(self) -> Dict[str, Dict[int, str]]:
+        if self.task == "image-segmentation":
+            return {
+                "logits": {0: "batch_size", 1: "num_queries"},
+                "pred_masks": {0: "batch_size", 1: "num_queries"},
+            }
+        else:
+            return super().outputs
 
 
 class YolosOnnxConfig(ViTOnnxConfig):
diff --git a/optimum/exporters/tasks.py b/optimum/exporters/tasks.py
--- a/optimum/exporters/tasks.py
+++ b/optimum/exporters/tasks.py
@@ -161,7 +161,7 @@ class TasksManager:
             "object-detection": "AutoModelForObjectDetection",
             "question-answering": "AutoModelForQuestionAnswering",
             "image-classification": "AutoModelForImageClassification",
-            "image-segmentation": "AutoModelForImageSegmentation",
+            "image-segmentation": ("AutoModelForImageSegmentation", "AutoModelForSemanticSegmentation"),
             "mask-generation": "AutoModel",
             "masked-im": "AutoModelForMaskedImageModeling",
             "semantic-segmentation": "AutoModelForSemanticSegmentation",
@@ -632,6 +632,7 @@ class TasksManager:
         "mobilevit": supported_tasks_mapping(
             "feature-extraction",
             "image-classification",
+            "image-segmentation",
             onnx="MobileViTOnnxConfig",
         ),
         "mobilenet-v1": supported_tasks_mapping(
@@ -766,6 +767,7 @@ class TasksManager:
         "segformer": supported_tasks_mapping(
             "feature-extraction",
             "image-classification",
+            "image-segmentation",
             "semantic-segmentation",
             onnx="SegformerOnnxConfig",
         ),


Test patch: