app.py

import os
import shutil
import torch
import gradio as gr
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,
)


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")

    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)
    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 = _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, cache_dir=folder)
            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:
                            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)

DESCRIPTION = """
This Space uses [Optimum Intel](https://huggingface.co/docs/optimum/main/en/intel/openvino/export) to automatically export a model from the [Hub](https://huggingface.co/models) to the [OpenVINO format](https://docs.openvino.ai/2024/documentation/openvino-ir-format.html).

The resulting model will then be pushed under your HF user namespace. For now we only support conversion for models that are hosted on public repositories, support for gated or private models will be supported in a future version of this space.
Custom models that needs to be loaded with `trust_remote_code=True` are also not supported.
"""

model_id = HuggingfaceHubSearch(
    label="Hub Model ID",
    placeholder="Search for model id on the hub",
    search_type="model",
)
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,
        private_repo,
    ],
    outputs=[
        gr.Markdown(label="output"),
    ],
    title="Export your model to OpenVINO",
    description=DESCRIPTION,
    api_name=False,
)

with gr.Blocks() as demo:
    gr.Markdown("You must be logged in to use this space")
    gr.LoginButton(min_width=250)
    interface.render()

demo.launch()