diff --git a/optimum/intel/utils/modeling_utils.py b/optimum/intel/utils/modeling_utils.py index 9b68266d16..1416556dba 100644 --- a/optimum/intel/utils/modeling_utils.py +++ b/optimum/intel/utils/modeling_utils.py @@ -12,6 +12,7 @@ # 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 @@ -19,6 +20,8 @@ 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] diff --git a/tests/openvino/test_modeling.py b/tests/openvino/test_modeling.py index ca5c4d26f9..1b164d4d37 100644 --- a/tests/openvino/test_modeling.py +++ b/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): diff --git a/tests/openvino/test_stable_diffusion.py b/tests/openvino/test_stable_diffusion.py index d8e8e76030..7467871bff 100644 --- a/tests/openvino/test_stable_diffusion.py +++ b/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) diff --git a/tests/openvino/utils_tests.py b/tests/openvino/utils_tests.py index 1f9b051d7d..79788e27b3 100644 --- a/tests/openvino/utils_tests.py +++ b/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",