Push new repo instead of PR

app.py

@@ -1,104 +1,166 @@
-import csv
 import os
-from datetime import datetime
-from typing import Optional, Union
+import shutil
+import torch
 import gradio as gr
-from huggingface_hub import HfApi, Repository
-from export import convert
+from huggingface_hub import HfApi, whoami, ModelCard
 from gradio_huggingfacehub_search import HuggingfaceHubSearch
+from textwrap import dedent
+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
+from optimum.intel import (
+    OVModelForAudioClassification,
+    OVModelForCausalLM,
+    OVModelForFeatureExtraction,
+    OVModelForImageClassification,
+    OVModelForMaskedLM,
+    OVModelForQuestionAnswering,
+    OVModelForSeq2SeqLM,
+    OVModelForSequenceClassification,
+    OVModelForTokenClassification,
+    OVStableDiffusionPipeline,
+    OVStableDiffusionXLPipeline,
+    OVLatentConsistencyModelPipeline,
+    OVModelForPix2Struct,
+    OVWeightQuantizationConfig,
+)
+from optimum.intel.openvino.modeling_diffusion import OVStableDiffusionPipelineBase
 
 
-DATASET_REPO_URL = "https://huggingface.co/datasets/optimum/exporters"
-DATA_FILENAME = "data.csv"
-DATA_FILE = os.path.join("openvino", DATA_FILENAME)
-HF_TOKEN = os.environ.get("HF_WRITE_TOKEN")
-DATA_DIR = "exporters_data"
+def export(
+    model_id: str,
+    private_repo: bool,
+    oauth_token: gr.OAuthToken,
+):
+    if oauth_token.token is None:
+        raise ValueError("You must be logged in to use this space")
 
-repo = None
-# if HF_TOKEN:
-#     repo = Repository(local_dir=DATA_DIR, clone_from=DATASET_REPO_URL, token=HF_TOKEN)
+    model_name = model_id.split("/")[-1]
+    username = whoami(oauth_token.token)["name"]
+    new_repo_id = f"{username}/{model_name}-openvino"
 
-def export(model_id: str, task: str, oauth_token: gr.OAuthToken) -> str:
+    task = TasksManager.infer_task_from_model(model_id)
+    if task not in _HEAD_TO_AUTOMODELS:
+        raise ValueError(
+            f"The task '{task}' is not supported, only {_HEAD_TO_AUTOMODELS.keys()} tasks are supported"
+        )
 
-    if oauth_token.token is None:
-        raise ValueError("You must be logged in to use this space")
-    token = oauth_token.token
-
-    if model_id == "" or token == "":
-        return """
-        ### Invalid input 🐞
-        Please fill a token and model name.
-        """
-    try:
-        api = HfApi(token=token)
-
-        error, commit_info = convert(api=api, model_id=model_id, task=task, force=False)
-        if error != "0":
-            return error
-
-        print("[commit_info]", commit_info)
-
-        # save in a private dataset
-        if repo is not None:
-            repo.git_pull(rebase=True)
-            with open(os.path.join(DATA_DIR, DATA_FILE), "a") as csvfile:
-                writer = csv.DictWriter(csvfile, fieldnames=["model_id", "pr_url", "time"])
-                writer.writerow(
-                    {
-                        "model_id": model_id,
-                        "pr_url": commit_info.pr_url,
-                        "time": str(datetime.now()),
-                    }
-                )
-            commit_url = repo.push_to_hub()
-            print("[dataset]", commit_url)
-
-        return f"#### Success 🔥 Yay! This model was successfully exported and a PR was open using your token, here: [{commit_info.pr_url}]({commit_info.pr_url})"
-    except Exception as e:
-        return f"#### Error: {e}"
-
-
-TTILE_IMAGE = """
-<div
-    style="
-        display: block;
-        margin-left: auto;
-        margin-right: auto;
-        width: 50%;
-    "
->
-<img src="https://huggingface.co/spaces/echarlaix/openvino-export/resolve/main/header.png"/>
-</div>
-"""
-
-TITLE = """
-<div
-    style="
-        display: inline-flex;
-        align-items: center;
-        text-align: center;
-        max-width: 1400px;
-        gap: 0.8rem;
-        font-size: 2.2rem;
-    "
->
-<h1 style="text-align:center; font-weight: 1200">
-    Export your model to OpenVINO
-</h1>
-</div>
-"""
-
-DESCRIPTION = """
-This Space uses [Optimum Intel](https://huggingface.co/docs/optimum/intel/inference) to automatically export your model to the OpenVINO format.
-
-After the model conversion, we will open a PR against the source repo to add the resulting model.
-
-To export your model you need:
-- A Model ID from the Hub
-
-That's it ! 🔥
-"""
+    if task == "text2text-generation":
+        raise ValueError("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:
+        raise Exception(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)
+            ov_model.save_pretrained(folder)
+
+            """
+            if not isinstance(ov_model, OVStableDiffusionPipelineBase):
+                model = TasksManager.get_model_from_task(task, model_id)
+                exporter_config_class = TasksManager.get_exporter_config_constructor(
+                    exporter="openvino",
+                    model=model,
+                    task=task,
+                    model_name=model_id,
+                    model_type=model.config.model_type.replace("_", "-"),
+                )
+                openvino_config = exporter_config_class(model.config)
+                inputs = openvino_config.generate_dummy_inputs(framework="pt")
+                ov_outputs = ov_model(**inputs)
+                outputs = model(**inputs)
+
+                for output_name in ov_outputs:
+                    if isinstance(outputs, torch.Tensor) and not torch.allclose(outputs[output_name], ov_outputs[output_name], atol=1e-3):
+                        raise ValueError(
+                            "The exported model does not have the same outputs as the original model. Export interrupted."
+                        )
+            """
+
+            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)
+            folder_parts = len(folder.parts)
+            for file_path in folder.glob("**/*"):
+                name = Path(*file_path.parts[folder_parts:])
+                if not file_path.is_file() or any(part_name.startswith(".") for part_name in name.parts):
+                    continue
+                try:
+                    api.upload_file(
+                        path_or_fileobj=file_path,
+                        path_in_repo=str(name),
+                        repo_id=new_repo_id,
+                    )
+                except Exception as e:
+                    raise Exception(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)
 
 
 model_id = HuggingfaceHubSearch(
@@ -106,21 +168,22 @@ model_id = HuggingfaceHubSearch(
     placeholder="Search for model id on the hub",
     search_type="model",
 )
-task = gr.Textbox(
-    value="auto",
-    label="Task : can be left to auto, will be automatically inferred",
+private_repo = gr.Checkbox(
+    value=False,
+    label="Private Repo",
+    info="Create a private repo under your username",
 )
 interface = gr.Interface(
     fn=export,
     inputs=[
         model_id,
-        task,
+        private_repo,
     ],
     outputs=[
         gr.Markdown(label="output"),
     ],
-    title=TITLE,
-    description=DESCRIPTION,
+    title="Export your model to OpenVINO",
+    description="This space converts your model to the OpenVINO format using [optimum-intel](https://huggingface.co/docs/optimum/main/intel/openvino/inference) The resulting model will then be pushed on the Hub under your HF user namespace",
     api_name=False,
 )
 

export.py

@@ -1,170 +0,0 @@
-import argparse
-import os
-import shutil
-from pathlib import Path
-from tempfile import TemporaryDirectory
-from typing import List, Optional, Tuple
-import torch
-
-from huggingface_hub import (
-    CommitInfo,
-    CommitOperationAdd,
-    Discussion,
-    HfApi,
-    get_repo_discussions,
-    hf_hub_download,
-)
-from huggingface_hub.file_download import repo_folder_name
-from optimum.exporters.onnx import validate_model_outputs
-from optimum.exporters.tasks import TasksManager
-from transformers import AutoConfig, AutoTokenizer, is_torch_available
-from optimum.intel.openvino import (
-    OVModelForAudioClassification,
-    OVModelForCausalLM,
-    OVModelForFeatureExtraction,
-    OVModelForImageClassification,
-    OVModelForMaskedLM,
-    OVModelForQuestionAnswering,
-    OVModelForSeq2SeqLM,
-    OVModelForSequenceClassification,
-    OVModelForTokenClassification,
-    OVStableDiffusionPipeline,
-)
-from optimum.intel.utils.constant import _TASK_ALIASES
-from optimum.intel.openvino.utils import _HEAD_TO_AUTOMODELS
-from optimum.exporters import TasksManager
-
-SPACES_URL = "https://huggingface.co/spaces/echarlaix/openvino-export"
-
-
-def previous_pr(api: "HfApi", model_id: str, pr_title: str) -> Optional["Discussion"]:
-    try:
-        discussions = api.get_repo_discussions(repo_id=model_id)
-    except Exception:
-        return None
-    for discussion in discussions:
-        if (
-            discussion.status == "open"
-            and discussion.is_pull_request
-            and discussion.title == pr_title
-        ):
-            return discussion
-
-
-def convert_openvino(model_id: str, task: str, folder: str) -> List:
-    task = TasksManager.map_from_synonym(task)
-    if task == "auto":
-        try:
-            task = TasksManager.infer_task_from_model(model_id)
-        except KeyError as e:
-            raise KeyError(
-                    f"The task could not be automatically inferred. Please provide the task argument with the relevant task from {', '.join(TasksManager.get_all_tasks())}.  {e}"
-            )
-
-    task = _TASK_ALIASES.get(task, task)
-    if task not in _HEAD_TO_AUTOMODELS:
-        raise ValueError(f"The task '{task}' is not supported, only {_HEAD_TO_AUTOMODELS.keys()} tasks are supported")
-
-    if task == "text2text-generation":
-        raise ValueError("Export of Seq2Seq models is currently disabled.")
-
-    auto_model_class = eval(_HEAD_TO_AUTOMODELS[task])
-    ov_model = auto_model_class.from_pretrained(model_id, export=True)
-    ov_model.save_pretrained(folder)
-    if not isinstance(ov_model, OVStableDiffusionPipeline):
-        try:
-            model = TasksManager.get_model_from_task(task, model_id)
-            exporter_config_class = TasksManager.get_exporter_config_constructor(
-                exporter="openvino",
-                model=model,
-                task=task,
-                model_name=model_id,
-                model_type=model.config.model_type.replace("_", "-"),
-            )
-            openvino_config = exporter_config_class(model.config)
-            inputs = openvino_config.generate_dummy_inputs(framework="pt")
-            ov_outputs = ov_model(**inputs)
-            outputs = model(**inputs)
-
-            for output_name in ov_outputs:
-                if isinstance(outputs, torch.Tensor) and not torch.allclose(outputs[output_name], ov_outputs[output_name], atol=1e-3):
-                    raise ValueError(
-                        "The exported model does not have the same outputs as the original model. Export interrupted."
-                    )
-        except Exception as e:
-            raise
-
-    file_names = {elem for elem in os.listdir(folder) if os.path.isfile(os.path.join(folder, elem))}
-
-    operations = [
-        CommitOperationAdd(
-            path_in_repo=file_name, path_or_fileobj=os.path.join(folder, file_name)
-        )
-        for file_name in file_names if "openvino" in file_name
-    ]
-
-    dir_names = set(os.listdir(folder)) - file_names
-
-    for dir_name in dir_names.intersection({"vae_encoder", "vae_decoder", "text_encoder", "unet"}):
-        operations += [
-            CommitOperationAdd(
-                path_in_repo=os.path.join(dir_name, file_name),
-                path_or_fileobj=os.path.join(folder, dir_name, file_name),
-            )
-            for file_name in os.listdir(os.path.join(folder, dir_name)) if "openvino" in file_name
-        ]
-
-    return operations
-
-
-def convert(
-    api: "HfApi",
-    model_id: str,
-    task: str,
-    force: bool = False,
-) -> Tuple[int, "CommitInfo"]:
-    pr_title = "Adding OpenVINO file of this model"
-    info = api.model_info(model_id)
-    filenames = set(s.rfilename for s in info.siblings)
-
-    requesting_user = api.whoami()["name"]
-
-    if task == "auto":
-        try:
-            task = TasksManager.infer_task_from_model(model_id)
-        except Exception as e:
-            return (
-                f"### Error: {e}. Please pass explicitely the task as it could not be infered.",
-                None,
-            )
-
-    with TemporaryDirectory() as d:
-        folder = os.path.join(d, repo_folder_name(repo_id=model_id, repo_type="models"))
-        os.makedirs(folder)
-        new_pr = None
-        try:
-            pr = previous_pr(api, model_id, pr_title)
-            if "openvino_model.xml" in filenames and not force:
-                raise Exception(f"Model {model_id} is already converted, skipping..")
-            elif pr is not None and not force:
-                url = f"https://huggingface.co/{model_id}/discussions/{pr.num}"
-                new_pr = pr
-                raise Exception(
-                    f"Model {model_id} already has an open PR check out [{url}]({url})"
-                )
-            else:
-                operations = convert_openvino(model_id, task, folder)
-
-                commit_description = f"""
-                Beep boop I am the [OpenVINO exporter bot 🤖]({SPACES_URL}). On behalf of [{requesting_user}](https://huggingface.co/{requesting_user}), I would like to add to this repository the exported OpenVINO model.
-                """
-                new_pr = api.create_commit(
-                    repo_id=model_id,
-                    operations=operations,
-                    commit_message=pr_title,
-                    commit_description=commit_description,
-                    create_pr=True,
-                )
-        finally:
-            shutil.rmtree(folder)
-        return "0", new_pr