simplify model loading

app.py

@@ -24,11 +24,9 @@ from optimum.intel import (
     OVModelForSeq2SeqLM,
     OVModelForSequenceClassification,
     OVModelForTokenClassification,
-    OVStableDiffusionPipeline,
-    OVStableDiffusionXLPipeline,
-    OVLatentConsistencyModelPipeline,
     OVModelForPix2Struct,
     OVWeightQuantizationConfig,
+    OVDiffusionPipeline,
 )
 from diffusers import ConfigMixin
 
@@ -41,9 +39,6 @@ _HEAD_TO_AUTOMODELS = {
     "question-answering": "OVModelForQuestionAnswering",
     "image-classification": "OVModelForImageClassification",
     "audio-classification": "OVModelForAudioClassification",
-    "stable-diffusion": "OVStableDiffusionPipeline",
-    "stable-diffusion-xl": "OVStableDiffusionXLPipeline",
-    "latent-consistency": "OVLatentConsistencyModelPipeline",
 }
 
 
@@ -61,27 +56,22 @@ def export(model_id: str, private_repo: bool, overwritte: bool, oauth_token: gr.
         library_name = TasksManager.infer_library_from_model(model_id, token=oauth_token.token)
 
         if library_name == "diffusers":
-            ConfigMixin.config_name = "model_index.json"
-            class_name = ConfigMixin.load_config(model_id, token=oauth_token.token)["_class_name"].lower()
-            if "xl" in class_name:
-                task = "stable-diffusion-xl"
-            elif "consistency" in class_name:
-                task = "latent-consistency"
-            else:
-                task = "stable-diffusion"
+            auto_model_class = "OVDiffusionPipeline"
         else:
             task = TasksManager.infer_task_from_model(model_id, token=oauth_token.token)
 
-        if task == "text2text-generation":
-            return "Export of Seq2Seq models is currently disabled"
+            if task == "text2text-generation":
+                return "Export of Seq2Seq models is currently disabled"
+
+            if task not in _HEAD_TO_AUTOMODELS:
+                return f"The task '{task}' is not supported, only {_HEAD_TO_AUTOMODELS.keys()} tasks are supported"
+
+            auto_model_class = _HEAD_TO_AUTOMODELS[task]
 
-        if task not in _HEAD_TO_AUTOMODELS:
-            return f"The task '{task}' is not supported, only {_HEAD_TO_AUTOMODELS.keys()} tasks are supported"
 
-        auto_model_class = _HEAD_TO_AUTOMODELS[task]
         ov_files = _find_files_matching_pattern(
             model_id,
-            pattern=r"(.*)?openvino(.*)?\_model.xml",
+            pattern=r"(.*)?openvino(.*)?\_model(.*)?.xml$",
             use_auth_token=oauth_token.token,
         )