fix(exporters/tasks): support local timm model paths in ONNX export

optimum/exporters/tasks.py

@@ -1174,7 +1174,12 @@ def get_model_from_task(
             )
 
         if library_name == "timm":
-            model = model_class(f"hf_hub:{model_name_or_path}", pretrained=True, exportable=True)
+            import os
+
+            if os.path.isdir(model_name_or_path):
+                model = model_class(model_name_or_path, pretrained=True, exportable=True)
+            else:
+                model = model_class(f"hf_hub:{model_name_or_path}", pretrained=True, exportable=True)
             model = model.to(torch_dtype).to(device)
         elif library_name == "sentence_transformers":
             token = model_kwargs.pop("token", None)

tests/exporters/common/test_tasks_manager.py

@@ -16,6 +16,7 @@
 import inspect
 from typing import Optional, Set
 from unittest import TestCase
+from unittest.mock import MagicMock, patch
 
 import pytest
 from transformers import BertConfig, Pix2StructForConditionalGeneration, VisualBertForQuestionAnswering
@@ -175,6 +176,39 @@ def test_custom_class(self):
         model = TasksManager.get_model_from_task("question-answering", "uclanlp/visualbert-vqa")
         self.assertTrue(isinstance(model, VisualBertForQuestionAnswering))
 
+    def test_get_model_from_task_timm_local_dir_skips_hf_hub_prefix(self):
+        # Regression test for local timm checkpoint loading: a local directory must be
+        # passed to the timm model constructor as-is, while a Hub id keeps the
+        # "hf_hub:" prefix. Both paths are exercised without network access by mocking
+        # the resolved model class and os.path.isdir.
+        created = MagicMock()
+        created.to.return_value = created
+        model_class = MagicMock(return_value=created)
+
+        with patch.object(TasksManager, "get_model_class_for_task", return_value=model_class):
+            with patch("os.path.isdir", return_value=True):
+                TasksManager.get_model_from_task(
+                    "image-classification",
+                    "/local/path/to/timm_model",
+                    framework="pt",
+                    library_name="timm",
+                )
+            model_class.assert_called_once_with(
+                "/local/path/to/timm_model", pretrained=True, exportable=True
+            )
+
+            model_class.reset_mock()
+            with patch("os.path.isdir", return_value=False):
+                TasksManager.get_model_from_task(
+                    "image-classification",
+                    "timm/mobilenetv3_large_100.ra_in1k",
+                    framework="pt",
+                    library_name="timm",
+                )
+            model_class.assert_called_once_with(
+                "hf_hub:timm/mobilenetv3_large_100.ra_in1k", pretrained=True, exportable=True
+            )
+
     def test_library_detection(self):
         self.assertEqual(
             TasksManager.infer_library_from_model("intfloat/multilingual-e5-large"), "sentence_transformers"