Timm Models integration to Optimum-intel

optimum/intel/openvino/modeling.py

@@ -11,14 +11,15 @@
 #  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
-
 import logging
+from pathlib import Path
 from typing import Optional, Union
 
 import numpy as np
 import openvino
 import torch
 import transformers
+from huggingface_hub import model_info
 from transformers import (
     AutoConfig,
     AutoModel,
@@ -31,6 +32,7 @@
     AutoModelForQuestionAnswering,
     AutoModelForSequenceClassification,
     AutoModelForTokenClassification,
+    PretrainedConfig,
 )
 from transformers.file_utils import add_start_docstrings, add_start_docstrings_to_model_forward
 from transformers.modeling_outputs import (
@@ -47,6 +49,7 @@
 from optimum.exporters import TasksManager
 
 from .modeling_base import OVBaseModel
+from .modeling_timm import TimmConfig, TimmForImageClassification, TimmOnnxConfig, is_timm_ov_dir
 
 
 logger = logging.getLogger(__name__)
@@ -497,6 +500,56 @@ class OVModelForImageClassification(OVModel):
     def __init__(self, model=None, config=None, **kwargs):
         super().__init__(model, config, **kwargs)
 
+    @classmethod
+    def from_pretrained(
+        cls,
+        model_id: Union[str, Path],
+        export: bool = False,
+        config: Optional["PretrainedConfig"] = None,
+        use_auth_token: Optional[Union[bool, str]] = None,
+        revision: Optional[str] = None,
+        force_download: bool = False,
+        cache_dir: Optional[str] = None,
+        subfolder: str = "",
+        local_files_only: bool = False,
+        task: Optional[str] = None,
+        trust_remote_code: bool = False,
+        **kwargs,
+    ):
+        # Fix the mismatch between timm_config and huggingface_config
+        local_timm_model = is_timm_ov_dir(model_id)
+        if local_timm_model or model_info(model_id).library_name == "timm":
+            config = TimmConfig.from_pretrained(model_id, **kwargs)
+            #  If locally saved timm model, dirrectly load
+            if local_timm_model:
+                return super()._from_pretrained(
+                    model_id=model_id,
+                    config=config,
+                )
+            model = TimmForImageClassification.from_pretrained(model_id, **kwargs)
+            onnx_config = TimmOnnxConfig(model.config)
+
+            return cls._to_onnx_to_load(
+                model=model,
+                config=config,
+                onnx_config=onnx_config,
+            )
+        else:
+            return super().from_pretrained(
+                model_id=model_id,
+                config=config,
+                export=export,
+                use_auth_token=use_auth_token,
+                revision=revision,
+                force_download=force_download,
+                cache_dir=cache_dir,
+                subfolder=subfolder,
+                local_files_only=local_files_only,
+                task=task,
+                trust_remote_code=trust_remote_code,
+                **kwargs,
+            )
+
     @add_start_docstrings_to_model_forward(
         IMAGE_INPUTS_DOCSTRING.format("batch_size, num_channels, height, width")
         + IMAGE_CLASSIFICATION_EXAMPLE.format(

optimum/intel/openvino/modeling_base.py

@@ -25,7 +25,7 @@
 from transformers import PretrainedConfig
 from transformers.file_utils import add_start_docstrings
 
-from optimum.exporters.onnx import export
+from optimum.exporters.onnx import OnnxConfig, export
 from optimum.exporters.tasks import TasksManager
 from optimum.modeling_base import OptimizedModel
 
@@ -276,6 +276,31 @@ def _from_transformers(
         )
 
         onnx_config = onnx_config_class(model.config)
+
+        return cls._to_onnx_to_load(
+            model=model,
+            config=config,
+            onnx_config=onnx_config,
+            use_auth_token=use_auth_token,
+            revision=revision,
+            force_download=force_download,
+            cache_dir=cache_dir,
+            local_files_only=local_files_only,
+        )
+
+    @classmethod
+    def _to_onnx_to_load(
+        cls,
+        model: PreTrainedModel,
+        config: PretrainedConfig,
+        onnx_config: OnnxConfig,
+        use_auth_token: Optional[Union[bool, str]] = None,
+        revision: Optional[str] = None,
+        force_download: bool = False,
+        cache_dir: Optional[str] = None,
+        local_files_only: bool = False,
+        **kwargs,
+    ):
         save_dir = TemporaryDirectory()
         save_dir_path = Path(save_dir.name)
 

optimum/intel/openvino/modeling_diffusion.py

@@ -217,9 +217,6 @@ def _from_pretrained(
                     cls.config_name,
                 }
             )
-            ignore_patterns = ["*.msgpack", "*.safetensors", "*pytorch_model.bin"]
-            if not from_onnx:
-                ignore_patterns.extend(["*.onnx", "*.onnx_data"])
             # Downloads all repo's files matching the allowed patterns
             model_id = snapshot_download(
                 model_id,
@@ -228,7 +225,7 @@ def _from_pretrained(
                 use_auth_token=use_auth_token,
                 revision=revision,
                 allow_patterns=allow_patterns,
-                ignore_patterns=ignore_patterns,
+                ignore_patterns=["*.msgpack", "*.safetensors", "*pytorch_model.bin"],
             )
         new_model_save_dir = Path(model_id)