diff --git a/optimum/intel/openvino/modeling_base.py b/optimum/intel/openvino/modeling_base.py index 7937deea52..a644b7adb7 100644 --- a/optimum/intel/openvino/modeling_base.py +++ b/optimum/intel/openvino/modeling_base.py @@ -20,16 +20,18 @@ from typing import Dict, Optional, Union import openvino -from huggingface_hub import hf_hub_download -from huggingface_hub.constants import HUGGINGFACE_HUB_CACHE +from huggingface_hub import HfApi, hf_hub_download, try_to_load_from_cache +from huggingface_hub.constants import HF_HUB_OFFLINE, HUGGINGFACE_HUB_CACHE from openvino import Core, convert_model from openvino._offline_transformations import apply_moc_transformations, compress_model_transformation -from transformers import GenerationConfig, PretrainedConfig +from transformers import AutoConfig, GenerationConfig, PretrainedConfig from transformers.file_utils import add_start_docstrings from transformers.generation import GenerationMixin from optimum.exporters.onnx import OnnxConfig -from optimum.modeling_base import OptimizedModel +from optimum.exporters.tasks import TasksManager +from optimum.modeling_base import FROM_PRETRAINED_START_DOCSTRING, OptimizedModel +from optimum.utils import CONFIG_NAME from ...exporters.openvino import export, main_export from ..utils.import_utils import is_nncf_available @@ -524,3 +526,156 @@ def can_generate(self) -> bool: if isinstance(self, GenerationMixin): return True return False + + @classmethod + @add_start_docstrings(FROM_PRETRAINED_START_DOCSTRING) + def from_pretrained( + cls, + model_id: Union[str, Path], + export: Optional[bool] = None, + force_download: bool = False, + use_auth_token: Optional[str] = None, + cache_dir: str = HUGGINGFACE_HUB_CACHE, + subfolder: str = "", + config: Optional[PretrainedConfig] = None, + local_files_only: bool = False, + trust_remote_code: bool = False, + revision: Optional[str] = None, + **kwargs, + ) -> "OptimizedModel": + """ + Returns: + `OptimizedModel`: The loaded optimized model. + """ + if isinstance(model_id, Path): + model_id = model_id.as_posix() + + from_transformers = kwargs.pop("from_transformers", None) + if from_transformers is not None: + logger.warning( + "The argument `from_transformers` is deprecated, and will be removed in optimum 2.0. Use `export` instead" + ) + export = from_transformers + + if len(model_id.split("@")) == 2: + if revision is not None: + logger.warning( + f"The argument `revision` was set to {revision} but will be ignored for {model_id.split('@')[1]}" + ) + model_id, revision = model_id.split("@") + + library_name = TasksManager.infer_library_from_model( + model_id, subfolder, revision, cache_dir, use_auth_token=use_auth_token + ) + + if library_name == "timm": + config = PretrainedConfig.from_pretrained(model_id, subfolder, revision) + + if config is None: + if os.path.isdir(os.path.join(model_id, subfolder)) and cls.config_name == CONFIG_NAME: + if CONFIG_NAME in os.listdir(os.path.join(model_id, subfolder)): + config = AutoConfig.from_pretrained( + os.path.join(model_id, subfolder, CONFIG_NAME), trust_remote_code=trust_remote_code + ) + elif CONFIG_NAME in os.listdir(model_id): + config = AutoConfig.from_pretrained( + os.path.join(model_id, CONFIG_NAME), trust_remote_code=trust_remote_code + ) + logger.info( + f"config.json not found in the specified subfolder {subfolder}. Using the top level config.json." + ) + else: + raise OSError(f"config.json not found in {model_id} local folder") + else: + config = cls._load_config( + model_id, + revision=revision, + cache_dir=cache_dir, + use_auth_token=use_auth_token, + force_download=force_download, + subfolder=subfolder, + trust_remote_code=trust_remote_code, + ) + elif isinstance(config, (str, os.PathLike)): + config = cls._load_config( + config, + revision=revision, + cache_dir=cache_dir, + use_auth_token=use_auth_token, + force_download=force_download, + subfolder=subfolder, + trust_remote_code=trust_remote_code, + ) + + if export is None: + export = cls._check_export_status( + model_id, revision, subfolder, cache_dir, local_files_only or HF_HUB_OFFLINE + ) + if not export and trust_remote_code: + logger.warning( + "The argument `trust_remote_code` is to be used along with export=True. It will be ignored." + ) + elif export and trust_remote_code is None: + trust_remote_code = False + + from_pretrained_method = cls._from_transformers if export else cls._from_pretrained + + return from_pretrained_method( + model_id=model_id, + config=config, + revision=revision, + cache_dir=cache_dir, + force_download=force_download, + use_auth_token=use_auth_token, + subfolder=subfolder, + local_files_only=local_files_only, + trust_remote_code=trust_remote_code, + **kwargs, + ) + + @classmethod + def _check_export_status( + cls, + model_id: Union[str, Path], + revision: Optional[str] = None, + subfolder: str = "", + cache_dir: str = HUGGINGFACE_HUB_CACHE, + local_files_only: bool = HF_HUB_OFFLINE, + ): + model_dir = Path(model_id) + if subfolder: + model_dir = model_dir / subfolder + if model_dir.is_dir(): + return ( + not (model_dir / OV_XML_FILE_NAME).exists() + or not (model_dir / OV_XML_FILE_NAME.replace(".xml", ".bin")).exists() + ) + normalized_subfolder = None if not subfolder else Path(subfolder).as_posix() + ov_model_path = OV_XML_FILE_NAME if not subfolder else f"{normalized_subfolder}/{OV_XML_FILE_NAME}" + cache_model_path = try_to_load_from_cache( + model_id, ov_model_path, cache_dir=cache_dir, revision=revision or "main", repo_type="model" + ) + cache_bin_path = try_to_load_from_cache( + model_id, + ov_model_path.replace(".xml", ".bin"), + cache_dir=cache_dir, + revision=revision or "main", + repo_type="model", + ) + cache_status = cache_bin_path is not None and cache_model_path is not None + + if not cache_status and not local_files_only: + hf_api = HfApi() + try: + model_info = hf_api.model_info(model_id, revision=revision or "main") + model_files = [ + file.rfilename + for file in model_info.siblings + if normalized_subfolder is None or file.rfilename.startswith(normalized_subfolder) + ] + ov_model_path = OV_XML_FILE_NAME if not subfolder else f"{normalized_subfolder}/{OV_XML_FILE_NAME}" + return ov_model_path not in model_files or ov_model_path.replace(".xml", ".bin") not in model_files + except Exception: + return True + + return not cache_status diff --git a/optimum/intel/openvino/modeling_base_seq2seq.py b/optimum/intel/openvino/modeling_base_seq2seq.py index fb53f9b2e2..50dddd0c14 100644 --- a/optimum/intel/openvino/modeling_base_seq2seq.py +++ b/optimum/intel/openvino/modeling_base_seq2seq.py @@ -20,8 +20,8 @@ from typing import Dict, Optional, Union import openvino -from huggingface_hub import hf_hub_download -from huggingface_hub.constants import HUGGINGFACE_HUB_CACHE +from huggingface_hub import HfApi, hf_hub_download, try_to_load_from_cache +from huggingface_hub.constants import HF_HUB_OFFLINE, HUGGINGFACE_HUB_CACHE from openvino._offline_transformations import apply_moc_transformations, compress_model_transformation from transformers import GenerationConfig, PretrainedConfig from transformers.file_utils import add_start_docstrings @@ -362,3 +362,62 @@ def half(self): def forward(self, *args, **kwargs): raise NotImplementedError + + @classmethod + def _check_export_status( + cls, + model_id: Union[str, Path], + revision: Optional[str] = None, + subfolder: str = "", + cache_dir: str = HUGGINGFACE_HUB_CACHE, + local_files_only: bool = HF_HUB_OFFLINE, + ): + model_dir = Path(model_id) + if subfolder is not None: + model_dir = model_dir / subfolder + if model_dir.is_dir(): + encoder_exists = (model_dir / OV_ENCODER_NAME).exists() and ( + model_dir / OV_ENCODER_NAME.replace(".xml", ".bin") + ).exists() + decoder_exists = (model_dir / OV_DECODER_NAME).exists() and ( + model_dir / OV_DECODER_NAME.replace(".xml", ".bin") + ).exists() + return not encoder_exists or not decoder_exists + + cache_status = [] + normalized_subfolder = None if subfolder is None else Path(subfolder).as_posix() + + for model_name in [OV_ENCODER_NAME, OV_DECODER_NAME]: + ov_model_path = model_name if subfolder is None else f"{normalized_subfolder}/{model_name}" + cache_model_path = try_to_load_from_cache( + model_id, ov_model_path, cache_dir=cache_dir, revision=revision or "main", repo_type="model" + ) + cache_bin_path = try_to_load_from_cache( + model_id, + ov_model_path.replace(".xml", ".bin"), + cache_dir=cache_dir, + revision=revision or "main", + repo_type="model", + ) + + cache_status.append(cache_model_path is not None and cache_bin_path is not None) + + if not all(cache_status) and not local_files_only: + hf_api = HfApi() + try: + model_info = hf_api.model_info(model_id, revision=revision or "main") + normalized_subfolder = None if not subfolder else Path(subfolder).as_posix() + model_files = [ + file.rfilename + for file in model_info.siblings + if normalized_subfolder is None or file.rfilename.startswith(normalized_subfolder) + ] + for model_name in [OV_ENCODER_NAME, OV_DECODER_NAME]: + ov_model_path = model_name if not subfolder else f"{normalized_subfolder}/{model_name}" + if ov_model_path not in model_files or ov_model_path.replace(".xml", ".bin") not in model_files: + return True + return False + except Exception: + return True + + return not all(cache_status) diff --git a/optimum/intel/openvino/modeling_diffusion.py b/optimum/intel/openvino/modeling_diffusion.py index 1b880e736c..5671d40999 100644 --- a/optimum/intel/openvino/modeling_diffusion.py +++ b/optimum/intel/openvino/modeling_diffusion.py @@ -35,8 +35,8 @@ ) from diffusers.schedulers.scheduling_utils import SCHEDULER_CONFIG_NAME from diffusers.utils import CONFIG_NAME, is_invisible_watermark_available -from huggingface_hub import snapshot_download -from huggingface_hub.constants import HUGGINGFACE_HUB_CACHE +from huggingface_hub import HfApi, snapshot_download, try_to_load_from_cache +from huggingface_hub.constants import HF_HUB_OFFLINE, HUGGINGFACE_HUB_CACHE from openvino._offline_transformations import compress_model_transformation from openvino.runtime import Core from transformers import CLIPFeatureExtractor, CLIPTokenizer @@ -411,6 +411,65 @@ def _from_transformers( **kwargs, ) + @classmethod + def _check_export_status( + cls, + model_id: Union[str, Path], + revision: Optional[str] = None, + subfolder: str = "", + cache_dir: str = HUGGINGFACE_HUB_CACHE, + local_files_only: bool = HF_HUB_OFFLINE, + ): + sub_model_dirs = [DIFFUSION_MODEL_UNET_SUBFOLDER, DIFFUSION_MODEL_VAE_DECODER_SUBFOLDER] + + def check_model_part_status(model_id, subfolder, revision): + model_dir = Path(model_id) + if subfolder: + model_dir = model_dir / subfolder + if model_dir.is_dir(): + return ( + not (model_dir / OV_XML_FILE_NAME).exists() + or not (model_dir / OV_XML_FILE_NAME.replace(".xml", ".bin")).exists() + ) + + normalized_subfolder = None if not subfolder else Path(subfolder).as_posix() + ov_model_path = OV_XML_FILE_NAME if not subfolder else f"{normalized_subfolder}/{OV_XML_FILE_NAME}" + cache_model_path = try_to_load_from_cache( + model_id, ov_model_path, cache_dir=cache_dir, revision=revision or "main", repo_type="model" + ) + cache_bin_path = try_to_load_from_cache( + model_id, + ov_model_path.replace(".xml", ".bin"), + cache_dir=cache_dir, + revision=revision or "main", + repo_type="model", + ) + cache_status = cache_bin_path is not None and cache_model_path is not None + + if not cache_status and not local_files_only: + hf_api = HfApi() + try: + model_info = hf_api.model_info(model_id, revision=revision or "main") + model_files = [ + file.rfilename + for file in model_info.siblings + if normalized_subfolder is None or file.rfilename.startswith(normalized_subfolder) + ] + ov_model_path = OV_XML_FILE_NAME if not subfolder else f"{normalized_subfolder}/{OV_XML_FILE_NAME}" + return ov_model_path not in model_files or ov_model_path.replace(".xml", ".bin") not in model_files + except Exception: + return True + return not cache_status + + for sub_model_dir in sub_model_dirs: + sub_model = sub_model_dir if not subfolder else f"{Path(subfolder).as_posix()}/{sub_model_dir}" + + status = check_model_part_status(model_id, sub_model, revision) + if status: + return True + + return False + def to(self, device: str): if isinstance(device, str): self._device = device.upper()