Repository: huggingface/optimum

Issue:
[Feature request] Image-to-Image for super-resolution
### Feature request


Running
```bash
optimum-cli export onnx -m caidas/swin2SR-realworld-sr-x4-64-bsrgan-psnr out
```

Results in
```py
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 225, in run
    main_export(
  File "/usr/local/lib/python3.10/dist-packages/optimum/exporters/onnx/__main__.py", line 306, in main_export
    model = TasksManager.get_model_from_task(
  File "/usr/local/lib/python3.10/dist-packages/optimum/exporters/tasks.py", line 1481, in get_model_from_task
    if TasksManager._TASKS_TO_LIBRARY[task.replace("-with-past", "")] == "transformers":
KeyError: 'image-to-image'
```

meaning this task is not yet supported. I also tried selecting the `masked-im` task, however that results in
```py
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 225, in run
    main_export(
  File "/usr/local/lib/python3.10/dist-packages/optimum/exporters/onnx/__main__.py", line 306, in main_export
    model = TasksManager.get_model_from_task(
  File "/usr/local/lib/python3.10/dist-packages/optimum/exporters/tasks.py", line 1513, 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 558, in from_pretrained
    raise ValueError(
ValueError: Unrecognized configuration class <class 'transformers.models.swin2sr.configuration_swin2sr.Swin2SRConfig'> for this kind of AutoModel: AutoModelForMaskedImageModeling.
Model type should be one of DeiTConfig, FocalNetConfig, SwinConfig, Swinv2Config, ViTConfig.
```

See here for a list of available swin2sr models: https://huggingface.co/models?other=swin2sr

### Motivation

Requested [here](https://github.com/xenova/transformers.js/issues/138) by @josephrocca. With WebGPU support right around the corner, I'd like to have a variety of vision models ready to create demos for (on top of, for example, SAM which is now supported).

### Your contribution

I will add support to transformers.js when ready


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
@@ -152,6 +152,9 @@ class OnnxConfig(ExportConfig, ABC):
         "image-classification": OrderedDict({"logits": {0: "batch_size"}}),
         "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"}}),
+        "image-to-image": OrderedDict(
+            {"reconstruction": {0: "batch_size", 1: "num_channels", 2: "height", 3: "width"}}
+        ),
         "mask-generation": OrderedDict({"logits": {0: "batch_size"}}),
         "masked-im": OrderedDict(
             {"reconstruction" if check_if_transformers_greater("4.29.0") else "logits": {0: "batch_size"}}
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
@@ -770,6 +770,10 @@ class SwinOnnxConfig(ViTOnnxConfig):
     pass
 
 
+class Swin2srOnnxConfig(SwinOnnxConfig):
+    pass
+
+
 class PoolFormerOnnxConfig(ViTOnnxConfig):
     NORMALIZED_CONFIG_CLASS = NormalizedVisionConfig
     ATOL_FOR_VALIDATION = 2e-3
diff --git a/optimum/exporters/tasks.py b/optimum/exporters/tasks.py
--- a/optimum/exporters/tasks.py
+++ b/optimum/exporters/tasks.py
@@ -168,6 +168,7 @@ class TasksManager:
             "fill-mask": "AutoModelForMaskedLM",
             "image-classification": "AutoModelForImageClassification",
             "image-segmentation": ("AutoModelForImageSegmentation", "AutoModelForSemanticSegmentation"),
+            "image-to-image": "AutoModelForImageToImage",
             "image-to-text": "AutoModelForVision2Seq",
             "mask-generation": "AutoModel",
             "masked-im": "AutoModelForMaskedImageModeling",
@@ -884,6 +885,11 @@ class TasksManager:
             "masked-im",
             onnx="SwinOnnxConfig",
         ),
+        "swin2sr": supported_tasks_mapping(
+            "feature-extraction",
+            "image-to-image",
+            onnx="Swin2srOnnxConfig",
+        ),
         "t5": supported_tasks_mapping(
             "feature-extraction",
             "feature-extraction-with-past",


Test patch:
diff --git a/tests/exporters/exporters_utils.py b/tests/exporters/exporters_utils.py
--- a/tests/exporters/exporters_utils.py
+++ b/tests/exporters/exporters_utils.py
@@ -120,6 +120,7 @@
     "splinter": "hf-internal-testing/tiny-random-SplinterModel",
     "squeezebert": "hf-internal-testing/tiny-random-SqueezeBertModel",
     "swin": "hf-internal-testing/tiny-random-SwinModel",
+    "swin2sr": "hf-internal-testing/tiny-random-Swin2SRModel",
     "t5": "hf-internal-testing/tiny-random-t5",
     "vit": "hf-internal-testing/tiny-random-vit",
     "yolos": "hf-internal-testing/tiny-random-YolosModel",