From c73fb2feb837c06a3cf7f6acce0a8ad1a62fb6a3 Mon Sep 17 00:00:00 2001
From: baskrahmer <baskrahmer@gmail.com>
Date: Fri, 27 Oct 2023 15:31:30 +0000
Subject: [PATCH 1/5] Add ONNX export support for swin2SR models

---
 optimum/exporters/onnx/base.py | 3 +++
 optimum/exporters/tasks.py     | 5 +++++
 2 files changed, 8 insertions(+)

diff --git a/optimum/exporters/onnx/base.py b/optimum/exporters/onnx/base.py
index 1e5704e893..d6e9f69fbd 100644
--- 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/tasks.py b/optimum/exporters/tasks.py
index 2a0f9076ce..a495bfa1a4 100644
--- 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,10 @@ class TasksManager:
             "masked-im",
             onnx="SwinOnnxConfig",
         ),
+        "swin2sr": supported_tasks_mapping(
+            "image-to-image",
+            onnx="SwinOnnxConfig",
+        ),
         "t5": supported_tasks_mapping(
             "feature-extraction",
             "feature-extraction-with-past",

From 8fb8a36060278768a04ddfcd9d17a517bab1012a Mon Sep 17 00:00:00 2001
From: baskrahmer <baskrahmer@gmail.com>
Date: Fri, 27 Oct 2023 17:40:44 +0000
Subject: [PATCH 2/5] Add feature extraction task

---
 optimum/exporters/tasks.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/optimum/exporters/tasks.py b/optimum/exporters/tasks.py
index a495bfa1a4..dcce49297a 100644
--- a/optimum/exporters/tasks.py
+++ b/optimum/exporters/tasks.py
@@ -886,6 +886,7 @@ class TasksManager:
             onnx="SwinOnnxConfig",
         ),
         "swin2sr": supported_tasks_mapping(
+            "feature-extraction",
             "image-to-image",
             onnx="SwinOnnxConfig",
         ),

From 9409153ea3c720ea624a35b5cb09e29e7a5a1e51 Mon Sep 17 00:00:00 2001
From: baskrahmer <baskrahmer@gmail.com>
Date: Fri, 27 Oct 2023 17:40:59 +0000
Subject: [PATCH 3/5] Add testing model

---
 tests/exporters/exporters_utils.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tests/exporters/exporters_utils.py b/tests/exporters/exporters_utils.py
index 105a5a7d77..e0640b7657 100644
--- 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",

From e1c0b59461fb1fd145e3c6b4eda7bf2578299e8c Mon Sep 17 00:00:00 2001
From: baskrahmer <baskrahmer@gmail.com>
Date: Wed, 8 Nov 2023 17:09:09 +0000
Subject: [PATCH 4/5] Add Swin2srOnnxConfig class

---
 optimum/exporters/onnx/model_configs.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/optimum/exporters/onnx/model_configs.py b/optimum/exporters/onnx/model_configs.py
index e1461c2a0c..04e4c0b11f 100644
--- 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

From 308b0972fcfc066248d5754980670c515af582e7 Mon Sep 17 00:00:00 2001
From: fxmarty <9808326+fxmarty@users.noreply.github.com>
Date: Thu, 9 Nov 2023 23:32:46 +0900
Subject: [PATCH 5/5] Update optimum/exporters/tasks.py

---
 optimum/exporters/tasks.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/optimum/exporters/tasks.py b/optimum/exporters/tasks.py
index dcce49297a..c26dc98da7 100644
--- a/optimum/exporters/tasks.py
+++ b/optimum/exporters/tasks.py
@@ -888,7 +888,7 @@ class TasksManager:
         "swin2sr": supported_tasks_mapping(
             "feature-extraction",
             "image-to-image",
-            onnx="SwinOnnxConfig",
+            onnx="Swin2srOnnxConfig",
         ),
         "t5": supported_tasks_mapping(
             "feature-extraction",