add support of gated and private models

app.py

@@ -10,9 +10,7 @@ from pathlib import Path
 from tempfile import TemporaryDirectory
 
 from huggingface_hub.file_download import repo_folder_name
-from optimum.exporters.tasks import TasksManager
 from optimum.intel.utils.constant import _TASK_ALIASES
-from optimum.intel.openvino.utils import _HEAD_TO_AUTOMODELS
 from optimum.exporters import TasksManager
 
 from optimum.intel.utils.modeling_utils import _find_files_matching_pattern
@@ -32,119 +30,148 @@ from optimum.intel import (
     OVModelForPix2Struct,
     OVWeightQuantizationConfig,
 )
-
-
-def export(model_id: str, private_repo: bool, oauth_token: gr.OAuthToken):
+from diffusers import ConfigMixin
+
+_HEAD_TO_AUTOMODELS = {
+    "feature-extraction": "OVModelForFeatureExtraction",
+    "fill-mask": "OVModelForMaskedLM",
+    "text-generation": "OVModelForCausalLM",
+    "text-classification": "OVModelForSequenceClassification",
+    "token-classification": "OVModelForTokenClassification",
+    "question-answering": "OVModelForQuestionAnswering",
+    "image-classification": "OVModelForImageClassification",
+    "audio-classification": "OVModelForAudioClassification",
+    "stable-diffusion": "OVStableDiffusionPipeline",
+    "stable-diffusion-xl": "OVStableDiffusionXLPipeline",
+    "latent-consistency": "OVLatentConsistencyModelPipeline",
+}
+
+
+def export(model_id: str, private_repo: bool, overwritte: bool, oauth_token: gr.OAuthToken):
     if oauth_token.token is None:
         return "You must be logged in to use this space"
 
     if not model_id:
-        return "### Invalid input 🐞 Please specify a model name, got {model_id}"
-
-    model_name = model_id.split("/")[-1]
-    username = whoami(oauth_token.token)["name"]
-    new_repo_id = f"{username}/{model_name}-openvino"
-    # task = TasksManager.infer_task_from_model(model_id, token=oauth_token.token)
-    task = TasksManager.infer_task_from_model(model_id)
-
-    if task not in _HEAD_TO_AUTOMODELS:
-        return f"The task '{task}' is not supported, only {_HEAD_TO_AUTOMODELS.keys()} tasks are supported"
-
-    if task == "text2text-generation":
-        return "Export of Seq2Seq models is currently disabled"
-
-    auto_model_class = _HEAD_TO_AUTOMODELS[task]
-    ov_files = _find_files_matching_pattern(
-        model_id,
-        pattern=r"(.*)?openvino(.*)?\_model.xml",
-        use_auth_token=oauth_token.token,
-    )
-
-    if len(ov_files) > 0:
-        return f"Model {model_id} is already converted, skipping.."
-
-    api = HfApi(token=oauth_token.token)
-
-    with TemporaryDirectory() as d:
-        folder = os.path.join(d, repo_folder_name(repo_id=model_id, repo_type="models"))
-        os.makedirs(folder)
-        try:
-            api.snapshot_download(repo_id=model_id, local_dir=folder, allow_patterns=["*.json"])
-            ov_model = eval(auto_model_class).from_pretrained(model_id, export=True, cache_dir=folder, token=oauth_token.token)
-            ov_model.save_pretrained(folder)
-            new_repo_url = api.create_repo(repo_id=new_repo_id, exist_ok=True, private=private_repo)
-            new_repo_id = new_repo_url.repo_id
-            print("Repo created successfully!", new_repo_url)
-
-            folder = Path(folder)
-            for dir_name in (
-                "",
-                "vae_encoder",
-                "vae_decoder",
-                "text_encoder",
-                "text_encoder_2",
-                "unet",
-                "tokenizer",
-                "tokenizer_2",
-                "scheduler",
-                "feature_extractor",
-            ):
-                if not (folder / dir_name).is_dir():
-                    continue
-                for file_path in (folder / dir_name).iterdir():
-                    if file_path.is_file():
-                        try:
-                            api.upload_file(
-                                path_or_fileobj=file_path,
-                                path_in_repo=os.path.join(dir_name, file_path.name),
-                                repo_id=new_repo_id,
-                            )
-                        except Exception as e:
-                            return f"Error uploading file {file_path}: {e}"
-
+        return f"### Invalid input 🐞 Please specify a model name, got {model_id}"
+
+    try:
+        model_name = model_id.split("/")[-1]
+        username = whoami(oauth_token.token)["name"]
+        new_repo_id = f"{username}/{model_name}-openvino"
+        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"
+        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 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",
+            use_auth_token=oauth_token.token,
+        )
+
+        if len(ov_files) > 0:
+            return f"Model {model_id} is already converted, skipping.."
+
+        api = HfApi(token=oauth_token.token)
+        if api.repo_exists(new_repo_id) and not overwritte:
+            return f"Model {new_repo_id} already exist, please set overwritte=True to push on an existing repo"
+
+        with TemporaryDirectory() as d:
+            folder = os.path.join(d, repo_folder_name(repo_id=model_id, repo_type="models"))
+            os.makedirs(folder)
             try:
-                card = ModelCard.load(model_id, token=oauth_token.token)
-            except:
-                card = ModelCard("")
-
-            if card.data.tags is None:
-                card.data.tags = []
-            card.data.tags.append("openvino")
-            card.data.base_model = model_id
-            card.text = dedent(
-                f"""
-                This model was converted to OpenVINO from [`{model_id}`](https://huggingface.co/{model_id}) using [optimum-intel](https://github.com/huggingface/optimum-intel)
-                via the [export](https://huggingface.co/spaces/echarlaix/openvino-export) space.
-
-                First make sure you have optimum-intel installed:
-
-                ```bash
-                pip install optimum[openvino]
-                ```
-
-                To load your model you can do as follows:
-
-                ```python
-                from optimum.intel import {auto_model_class}
-
-                model_id = "{new_repo_id}"
-                model = {auto_model_class}.from_pretrained(model_id)
-                ```
-                """
-            )
-            card_path = os.path.join(folder, "README.md")
-            card.save(card_path)
-
-            api.upload_file(
-                path_or_fileobj=card_path,
-                path_in_repo="README.md",
-                repo_id=new_repo_id,
-            )
-            return f"This model was successfully exported, find it under your repo {new_repo_url}'"
-        except Exception as e:
-            return f"### Error: {e}"
-        finally:
-            shutil.rmtree(folder, ignore_errors=True)
+                api.snapshot_download(repo_id=model_id, local_dir=folder, allow_patterns=["*.json"])
+                ov_model = eval(auto_model_class).from_pretrained(model_id, export=True, cache_dir=folder, token=oauth_token.token)
+                ov_model.save_pretrained(folder)
+                new_repo_url = api.create_repo(repo_id=new_repo_id, exist_ok=True, private=private_repo)
+                new_repo_id = new_repo_url.repo_id
+                print("Repo created successfully!", new_repo_url)
+
+                folder = Path(folder)
+                for dir_name in (
+                    "",
+                    "vae_encoder",
+                    "vae_decoder",
+                    "text_encoder",
+                    "text_encoder_2",
+                    "unet",
+                    "tokenizer",
+                    "tokenizer_2",
+                    "scheduler",
+                    "feature_extractor",
+                ):
+                    if not (folder / dir_name).is_dir():
+                        continue
+                    for file_path in (folder / dir_name).iterdir():
+                        if file_path.is_file():
+                            try:
+                                api.upload_file(
+                                    path_or_fileobj=file_path,
+                                    path_in_repo=os.path.join(dir_name, file_path.name),
+                                    repo_id=new_repo_id,
+                                )
+                            except Exception as e:
+                                return f"Error uploading file {file_path}: {e}"
+
+                try:
+                    card = ModelCard.load(model_id, token=oauth_token.token)
+                except:
+                    card = ModelCard("")
+
+                if card.data.tags is None:
+                    card.data.tags = []
+                card.data.tags.append("openvino")
+                card.data.base_model = model_id
+                card.text = dedent(
+                    f"""
+                    This model was converted to OpenVINO from [`{model_id}`](https://huggingface.co/{model_id}) using [optimum-intel](https://github.com/huggingface/optimum-intel)
+                    via the [export](https://huggingface.co/spaces/echarlaix/openvino-export) space.
+
+                    First make sure you have optimum-intel installed:
+
+                    ```bash
+                    pip install optimum[openvino]
+                    ```
+
+                    To load your model you can do as follows:
+
+                    ```python
+                    from optimum.intel import {auto_model_class}
+
+                    model_id = "{new_repo_id}"
+                    model = {auto_model_class}.from_pretrained(model_id)
+                    ```
+                    """
+                )
+                card_path = os.path.join(folder, "README.md")
+                card.save(card_path)
+
+                api.upload_file(
+                    path_or_fileobj=card_path,
+                    path_in_repo="README.md",
+                    repo_id=new_repo_id,
+                )
+                return f"This model was successfully exported, find it under your repo {new_repo_url}'"
+            finally:
+                shutil.rmtree(folder, ignore_errors=True)
+    except Exception as e:
+        return f"### Error: {e}"
 
 DESCRIPTION = """
 This Space uses [Optimum Intel](https://huggingface.co/docs/optimum/main/en/intel/openvino/export) to automatically export a model from the Hub to the [OpenVINO format](https://docs.openvino.ai/2024/documentation/openvino-ir-format.html).
@@ -164,11 +191,17 @@ private_repo = gr.Checkbox(
     label="Private Repo",
     info="Create a private repo under your username",
 )
+overwritte = gr.Checkbox(
+    value=False,
+    label="Overwrite repo content",
+    info="Push files on existing repo potentially overwriting existing files",
+)
 interface = gr.Interface(
     fn=export,
     inputs=[
         model_id,
         private_repo,
+        overwritte,
     ],
     outputs=[
         gr.Markdown(label="output"),

requirements.txt

@@ -3,7 +3,6 @@ gradio[oauth]>=4.37.2
 gradio_huggingfacehub_search==0.0.6
 transformers==4.42.4
 diffusers==0.29.1
-optimum==1.21.2
-optimum-intel==1.18.1
+git+https://github.com/huggingface/optimum-intel.git
 openvino
 nncf