add test

optimum/intel/utils/modeling_utils.py

@@ -12,13 +12,16 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 
+import os
 import re
 from pathlib import Path
 from typing import List, Optional, Union
 
 import torch
 from huggingface_hub import HfApi, HfFolder
 
+from optimum.exporters import TasksManager
+
 
 MULTI_QUERY_ATTN_MODELS = {"falcon", "gpt_bigcode"}
 
@@ -93,19 +96,26 @@ def _find_files_matching_pattern(
     Returns:
         `List[Path]`
     """
-    model_path = Path(model_name_or_path) if isinstance(model_name_or_path, str) else model_name_or_path
-    pattern = re.compile(f"{subfolder}/{pattern}" if subfolder != "" else pattern)
-    subfolder = subfolder or "."
+    model_path = Path(model_name_or_path) if not isinstance(model_name_or_path, Path) else model_name_or_path
+
+    if isinstance(use_auth_token, bool):
+        token = HfFolder().get_token()
+    else:
+        token = use_auth_token
+
+    library_name = TasksManager.infer_library_from_model(
+        model_name_or_path, subfolder=subfolder, revision=revision, token=token
+    )
+    if library_name == "diffusers":
+        subfolder = os.path.join(subfolder, "unet")
+    else:
+        subfolder = subfolder or "."
 
     if model_path.is_dir():
         glob_pattern = subfolder + "/*"
         files = model_path.glob(glob_pattern)
         files = [p for p in files if re.search(pattern, str(p))]
     else:
-        if isinstance(use_auth_token, bool):
-            token = HfFolder().get_token()
-        else:
-            token = use_auth_token
         repo_files = map(Path, HfApi().list_repo_files(model_name_or_path, revision=revision, token=token))
         files = [Path(p) for p in repo_files if re.match(pattern, str(p)) and str(p.parent) == subfolder]
 

tests/openvino/test_modeling.py

@@ -27,6 +27,7 @@
 import torch
 from datasets import load_dataset
 from evaluate import evaluator
+from huggingface_hub import HfApi
 from parameterized import parameterized
 from PIL import Image
 from transformers import (
@@ -84,6 +85,7 @@
 from optimum.intel.openvino.utils import _print_compiled_model_properties
 from optimum.intel.pipelines import pipeline as optimum_pipeline
 from optimum.intel.utils.import_utils import is_openvino_version, is_transformers_version
+from optimum.intel.utils.modeling_utils import _find_files_matching_pattern
 from optimum.utils import (
     DIFFUSION_MODEL_TEXT_ENCODER_SUBFOLDER,
     DIFFUSION_MODEL_UNET_SUBFOLDER,
@@ -275,6 +277,48 @@ def test_infer_export_when_loading(self):
         del model
         gc.collect()
 
+    def test_find_files_matching_pattern(self):
+        model_id = "echarlaix/tiny-random-PhiForCausalLM"
+        pattern = r"(.*)?openvino(.*)?\_model.xml"
+
+        # hub model
+        for revision in ("ov", ""):
+            ov_files = _find_files_matching_pattern(model_id, pattern=pattern, revision=revision)
+            self.assertTrue(len(ov_files) == 0 if revision == "" else len(ov_files) > 0)
+
+        ov_files = _find_files_matching_pattern(model_id, pattern=pattern, subfolder="openvino")
+        self.assertTrue(len(ov_files) > 0)
+
+        # local model
+        api = HfApi()
+        with tempfile.TemporaryDirectory() as tmpdirname:
+            for revision in ("ov", ""):
+                local_dir = Path(tmpdirname) / revision
+                api.snapshot_download(repo_id=model_id, local_dir=local_dir, revision=revision)
+
+                ov_files = _find_files_matching_pattern(local_dir, pattern=pattern, revision=revision)
+                self.assertTrue(len(ov_files) == 0 if revision == "" else len(ov_files) > 0)
+
+                if revision == "":
+                    ov_files = _find_files_matching_pattern(local_dir, pattern=pattern, subfolder="openvino")
+                    self.assertTrue(len(ov_files) > 0)
+
+    @parameterized.expand(("stable-diffusion", "stable-diffusion-openvino"))
+    def test_find_files_matching_pattern_sd(self, model_arch):
+        pattern = r"(.*)?openvino(.*)?\_model.xml"
+        model_id = MODEL_NAMES[model_arch]
+        # hub model
+        ov_files = _find_files_matching_pattern(model_id, pattern=pattern)
+        self.assertTrue(len(ov_files) > 0 if "openvino" in model_id else len(ov_files) == 0)
+
+        # local model
+        api = HfApi()
+        with tempfile.TemporaryDirectory() as tmpdirname:
+            local_dir = Path(tmpdirname) / "model"
+            api.snapshot_download(repo_id=model_id, local_dir=local_dir)
+            ov_files = _find_files_matching_pattern(local_dir, pattern=pattern)
+            self.assertTrue(len(ov_files) > 0 if "openvino" in model_id else len(ov_files) == 0)
+
 
 class PipelineTest(unittest.TestCase):
     def test_load_model_from_hub(self):

tests/openvino/test_stable_diffusion.py

@@ -94,7 +94,7 @@ class OVStableDiffusionPipelineBaseTest(unittest.TestCase):
     @parameterized.expand(SUPPORTED_ARCHITECTURES)
     def test_num_images_per_prompt(self, model_arch: str):
         model_id = MODEL_NAMES[model_arch]
-        pipeline = self.MODEL_CLASS.from_pretrained(model_id, export=True, compile=False)
+        pipeline = self.MODEL_CLASS.from_pretrained(model_id, compile=False)
         pipeline.to("cpu")
         pipeline.compile()
         self.assertEqual(pipeline.vae_scale_factor, 2)

tests/openvino/utils_tests.py

@@ -109,6 +109,7 @@
     "speech_to_text": "hf-internal-testing/tiny-random-Speech2TextModel",
     "squeezebert": "hf-internal-testing/tiny-random-squeezebert",
     "stable-diffusion": "hf-internal-testing/tiny-stable-diffusion-torch",
+    "stable-diffusion-openvino": "hf-internal-testing/tiny-stable-diffusion-openvino",
     "stable-diffusion-xl": "echarlaix/tiny-random-stable-diffusion-xl",
     "stable-diffusion-xl-refiner": "echarlaix/tiny-random-stable-diffusion-xl-refiner",
     "stablelm": "hf-internal-testing/tiny-random-StableLmForCausalLM",