OpenVINO/export
46
1import os2import shutil3import torch4import gradio as gr5from huggingface_hub import HfApi, whoami, ModelCard, model_info6from gradio_huggingfacehub_search import HuggingfaceHubSearch7from textwrap import dedent8from pathlib import Path9 10from tempfile import TemporaryDirectory11 12from huggingface_hub.file_download import repo_folder_name13from optimum.intel.utils.constant import _TASK_ALIASES14from optimum.exporters.tasks import TasksManager15 16from optimum.intel.utils.modeling_utils import _find_files_matching_pattern17from optimum.intel import (18 OVModelForAudioClassification,19 OVModelForCausalLM,20 OVModelForFeatureExtraction,21 OVModelForImageClassification,22 OVModelForMaskedLM,23 OVModelForQuestionAnswering,24 OVModelForSeq2SeqLM,25 OVModelForSequenceClassification,26 OVModelForTokenClassification,27 OVModelForPix2Struct,28 OVModelForVisualCausalLM,29 OVWeightQuantizationConfig,30 OVDiffusionPipeline,31)32 33_HEAD_TO_AUTOMODELS = {34 "feature-extraction": "OVModelForFeatureExtraction",35 "fill-mask": "OVModelForMaskedLM",36 "text-generation": "OVModelForCausalLM",37 "text-classification": "OVModelForSequenceClassification",38 "token-classification": "OVModelForTokenClassification",39 "question-answering": "OVModelForQuestionAnswering",40 "image-classification": "OVModelForImageClassification",41 "audio-classification": "OVModelForAudioClassification",42 "image-text-to-text": "OVModelForVisualCausalLM"43}44 45 46def export(model_id: str, private_repo: bool, overwritte: bool, oauth_token: gr.OAuthToken):47 if oauth_token.token is None:48 return "You must be logged in to use this space"49 50 if not model_id:51 return f"### Invalid input 🐞 Please specify a model name, got {model_id}"52 53 try:54 model_name = model_id.split("/")[-1]55 username = whoami(oauth_token.token)["name"]56 new_repo_id = f"{username}/{model_name}-openvino"57 library_name = TasksManager.infer_library_from_model(model_id, token=oauth_token.token)58 59 if library_name == "diffusers":60 auto_model_class = "OVDiffusionPipeline"61 elif library_name == "transformers":62 task = TasksManager.infer_task_from_model(model_id, token=oauth_token.token)63 64 if task == "text2text-generation":65 return "Export of Seq2Seq models is currently disabled"66 67 if task not in _HEAD_TO_AUTOMODELS:68 return f"The task '{task}' is not supported, only {_HEAD_TO_AUTOMODELS.keys()} tasks are supported"69 70 auto_model_class = _HEAD_TO_AUTOMODELS[task]71 else:72 # TODO: add sentence-transformers and timm support in space73 return f"Library {library_name} not yet supported"74 75 ov_files = _find_files_matching_pattern(76 model_id,77 pattern=r"(.*)?openvino(.*)?\_model(.*)?.xml$",78 use_auth_token=oauth_token.token,79 )80 81 if len(ov_files) > 0:82 return f"Model {model_id} is already converted, skipping.."83 84 api = HfApi(token=oauth_token.token)85 if api.repo_exists(new_repo_id) and not overwritte:86 return f"Model {new_repo_id} already exist, please tick the overwritte box to push on an existing repository"87 88 with TemporaryDirectory() as d:89 folder = os.path.join(d, repo_folder_name(repo_id=model_id, repo_type="models"))90 os.makedirs(folder)91 try:92 api.snapshot_download(repo_id=model_id, local_dir=folder, allow_patterns=["*.json"])93 ov_model = eval(auto_model_class).from_pretrained(model_id, export=True, cache_dir=folder, token=oauth_token.token)94 ov_model.save_pretrained(folder)95 new_repo_url = api.create_repo(repo_id=new_repo_id, exist_ok=True, private=private_repo)96 new_repo_id = new_repo_url.repo_id97 print("Repository created successfully!", new_repo_url)98 99 folder = Path(folder)100 for dir_name in (101 "",102 "vae_encoder",103 "vae_decoder",104 "text_encoder",105 "text_encoder_2",106 "unet",107 "tokenizer",108 "tokenizer_2",109 "scheduler",110 "feature_extractor",111 ):112 if not (folder / dir_name).is_dir():113 continue114 for file_path in (folder / dir_name).iterdir():115 if file_path.is_file():116 try:117 api.upload_file(118 path_or_fileobj=file_path,119 path_in_repo=os.path.join(dir_name, file_path.name),120 repo_id=new_repo_id,121 )122 except Exception as e:123 return f"Error uploading file {file_path}: {e}"124 125 try:126 card = ModelCard.load(model_id, token=oauth_token.token)127 except:128 card = ModelCard("")129 130 if card.data.tags is None:131 card.data.tags = []132 card.data.tags.append("openvino")133 card.data.tags.append("openvino-export")134 card.data.base_model = model_id135 136 pipeline_tag = getattr(model_info(model_id, token=oauth_token.token), "pipeline_tag", None)137 if pipeline_tag is not None:138 card.data.pipeline_tag = pipeline_tag139 140 card.text = dedent(141 f"""142 This model was converted to OpenVINO from [`{model_id}`](https://huggingface.co/{model_id}) using [optimum-intel](https://github.com/huggingface/optimum-intel)143 via the [export](https://huggingface.co/spaces/echarlaix/openvino-export) space.144 145 First make sure you have optimum-intel installed:146 147 ```bash148 pip install optimum-intel149 ```150 151 To load your model you can do as follows:152 153 ```python154 from optimum.intel import {auto_model_class}155 156 model_id = "{new_repo_id}"157 model = {auto_model_class}.from_pretrained(model_id)158 ```159 """160 )161 card_path = os.path.join(folder, "README.md")162 card.save(card_path)163 164 api.upload_file(165 path_or_fileobj=card_path,166 path_in_repo="README.md",167 repo_id=new_repo_id,168 )169 return f"This model was successfully exported, find it under your repository {new_repo_url}"170 finally:171 shutil.rmtree(folder, ignore_errors=True)172 except Exception as e:173 return f"### Error: {e}"174 175DESCRIPTION = """176This 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).177 178After conversion, a repository will be pushed under your namespace with the resulting model.179 180The list of supported architectures can be found in the [documentation](https://huggingface.co/docs/optimum/main/en/intel/openvino/models).181"""182 183model_id = HuggingfaceHubSearch(184 label="Hub Model ID",185 placeholder="Search for model ID on the hub",186 search_type="model",187)188private_repo = gr.Checkbox(189 value=False,190 label="Private repository",191 info="Create a private repository instead of a public one",192)193overwritte = gr.Checkbox(194 value=False,195 label="Overwrite repository content",196 info="Enable pushing files on existing repositories, potentially overwriting existing files",197)198interface = gr.Interface(199 fn=export,200 inputs=[201 model_id,202 private_repo,203 overwritte,204 ],205 outputs=[206 gr.Markdown(label="output"),207 ],208 title="Export your model to OpenVINO",209 description=DESCRIPTION,210 api_name=False,211)212 213with gr.Blocks() as demo:214 gr.Markdown("You must be logged in to use this space")215 gr.LoginButton(min_width=250)216 interface.render()217 218demo.launch()219