optionally enable export if not exported model provided

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

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)

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