fix unit tests

Author: philschmid

src/huggingface_inference_toolkit/diffusers_utils.py

@@ -1,6 +1,9 @@
 import importlib.util
 import logging
 
+logger = logging.getLogger(__name__)
+logging.basicConfig(format="%(asctime)s | %(levelname)s | %(message)s", level=logging.INFO)
+
 _diffusers = importlib.util.find_spec("diffusers") is not None
 
 
@@ -12,9 +15,6 @@ def is_diffusers_available():
     import torch
     from diffusers import AutoPipelineForText2Image, DPMSolverMultistepScheduler, StableDiffusionPipeline
 
-logger = logging.getLogger(__name__)
-logging.basicConfig(format="%(asctime)s | %(levelname)s | %(message)s", level=logging.INFO)
-
 
 class IEAutoPipelineForText2Image:
     def __init__(self, model_dir: str, device: str = None):  # needs "cuda" for GPU

src/huggingface_inference_toolkit/utils.py

@@ -59,7 +59,7 @@ def create_artifact_filter(framework):
     """
     Returns a list of regex pattern based on the DL Framework. which will be to used to ignore files when downloading
     """
-    ignore_regex_list = list(framework2weight.values())
+    ignore_regex_list = list(set(framework2weight.values()))
 
     pattern = framework2weight.get(framework, None)
     if pattern in ignore_regex_list:
@@ -153,7 +153,6 @@ def _load_repository_from_hf(
         if any(f.rfilename.endswith("safetensors") for f in files):
             framework = "safetensors"
 
-
     # create regex to only include the framework specific weights
     ignore_regex = create_artifact_filter(framework)
     logger.info(f"Ignore regex pattern for files, which are not downloaded: { ', '.join(ignore_regex) }")
@@ -284,8 +283,8 @@ def convert_params_to_int_or_bool(params):
     for k, v in params.items():
         if v.isnumeric():
             params[k] = int(v)
-        if v == 'false':
+        if v == "false":
             params[k] = False
-        if v == 'true':
+        if v == "true":
             params[k] = True
     return params

tests/unit/test_diffusers.py

@@ -3,8 +3,8 @@
 from PIL import Image
 from transformers.testing_utils import require_torch, slow
 
-from huggingface_inference_toolkit.handler import get_inference_handler_either_custom_or_default_handler
-from huggingface_inference_toolkit.diffusers_utils import get_diffusers_pipeline, DiffusersPipelineImageToText
+
+from huggingface_inference_toolkit.diffusers_utils import get_diffusers_pipeline, IEAutoPipelineForText2Image
 from huggingface_inference_toolkit.utils import _load_repository_from_hf, get_pipeline
 
 
@@ -15,7 +15,7 @@ def test_get_diffusers_pipeline():
             "hf-internal-testing/tiny-stable-diffusion-torch", tmpdirname, framework="pytorch"
         )
         pipe = get_pipeline("text-to-image", storage_dir.as_posix())
-        assert isinstance(pipe, DiffusersPipelineImageToText)
+        assert isinstance(pipe, IEAutoPipelineForText2Image)
 
 
 @slow