| 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.intel.utils.constant import _TASK_ALIASES |
| 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 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 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 tick the overwritte box to push on an existing repository" |
|
|
| 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("Repository 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.tags.append("openvino-export") |
| 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 repository {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 IR format](https://docs.openvino.ai/2024/documentation/openvino-ir-format.html). |
| |
| After conversion, a repository will be pushed under your namespace with the resulting model. |
| |
| The list of supported architectures can be found in the [documentation](https://huggingface.co/docs/optimum/main/en/intel/openvino/models). |
| """ |
|
|
| 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 repository", |
| info="Create a private repository instead of a public one", |
| ) |
| overwritte = gr.Checkbox( |
| value=False, |
| label="Overwrite repository content", |
| info="Enable pushing files on existing repositories, potentially overwriting existing files", |
| ) |
| interface = gr.Interface( |
| fn=export, |
| inputs=[ |
| model_id, |
| private_repo, |
| overwritte, |
| ], |
| 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() |
|
|
