From 093e7c3ff822d64064a58c50074beaa1df25f918 Mon Sep 17 00:00:00 2001
From: IlyasMoutawwakil <moutawwakil.ilyas.tsi@gmail.com>
Date: Mon, 4 Nov 2024 11:59:04 +0100
Subject: [PATCH] fix

---
 optimum/exporters/tasks.py | 20 ++++++++++----------
 optimum/exporters/utils.py |  2 +-
 2 files changed, 11 insertions(+), 11 deletions(-)

diff --git a/optimum/exporters/tasks.py b/optimum/exporters/tasks.py
index 5b78323013..b4bce4696f 100644
--- a/optimum/exporters/tasks.py
+++ b/optimum/exporters/tasks.py
@@ -335,23 +335,23 @@ class TasksManager:
     }
 
     _DIFFUSERS_SUPPORTED_MODEL_TYPE = {
-        "t5-encoder-model": supported_tasks_mapping(
+        "t5-encoder": supported_tasks_mapping(
             "feature-extraction",
             onnx="T5EncoderOnnxConfig",
         ),
-        "clip-text-model": supported_tasks_mapping(
+        "clip-text": supported_tasks_mapping(
             "feature-extraction",
             onnx="CLIPTextOnnxConfig",
         ),
-        "clip-text-model-with-projection": supported_tasks_mapping(
+        "clip-text-with-projection": supported_tasks_mapping(
             "feature-extraction",
             onnx="CLIPTextWithProjectionOnnxConfig",
         ),
-        "flux-transformer-2d-model": supported_tasks_mapping(
+        "flux-transformer-2d": supported_tasks_mapping(
             "semantic-segmentation",
             onnx="FluxTransformerOnnxConfig",
         ),
-        "sd3-transformer-2d-model": supported_tasks_mapping(
+        "sd3-transformer-2d": supported_tasks_mapping(
             "semantic-segmentation",
             onnx="SD3TransformerOnnxConfig",
         ),
@@ -1190,11 +1190,11 @@ class TasksManager:
     }
     _UNSUPPORTED_CLI_MODEL_TYPE = {
         # diffusers model types
-        "clip-text-model",
-        "clip-text-model-with-projection",
-        "flux-transformer-2d-model",
-        "sd3-transformer-2d-model",
-        "t5-encoder-model",
+        "clip-text",
+        "clip-text-with-projection",
+        "flux-transformer-2d",
+        "sd3-transformer-2d",
+        "t5-encoder",
         "unet-2d-condition",
         "vae-encoder",
         "vae-decoder",
diff --git a/optimum/exporters/utils.py b/optimum/exporters/utils.py
index 3377b591b8..60de169de5 100644
--- a/optimum/exporters/utils.py
+++ b/optimum/exporters/utils.py
@@ -75,7 +75,7 @@
 
 
 _DIFFUSERS_CLASS_NAME_TO_SUBMODEL_TYPE = {
-    "CLIPTextModel": "clip-text-model",
+    "CLIPTextModel": "clip-text",
     "CLIPTextModelWithProjection": "clip-text-with-projection",
     "FluxTransformer2DModel": "flux-transformer-2d",
     "SD3Transformer2DModel": "sd3-transformer-2d",