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__)
@@ -481,6 +484,20 @@ def forward(
     >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
     >>> outputs = pipe(url)
     ```
+    This class can also be used with [timm](https://github.com/huggingface/pytorch-image-models)
+    models hosted on [HuggingFaceHub](https://huggingface.co/timm). Example:
+    ```python
+    >>> from transformers import pipeline
+    >>> from optimum.intel.openvino.modeling_timm import TimmImageProcessor
+    >>> from optimum.intel import OVModelForImageClassification
+
+    >>> model_id = "timm/vit_tiny_patch16_224.augreg_in21k"
+    >>> preprocessor = TimmImageProcessor.from_pretrained(model_id)
+    >>> model = OVModelForImageClassification.from_pretrained(model_id, export=True)
+    >>> pipe = pipeline("image-classification", model=model, feature_extractor=preprocessor)
+    >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
+    >>> outputs = pipe(url)
+    ```
 """
 
 
@@ -497,6 +514,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)