updated and tested stuff

Author: philschmid

README.md

@@ -40,6 +40,8 @@ docker build -t starlette-transformers:gpu -f dockerfiles/tensorflow/gpu/Dockerf
 ```bash
 docker run -ti -p 5000:5000 -e HF_MODEL_ID=distilbert-base-uncased-distilled-squad -e HF_TASK=question-answering starlette-transformers:cpu
 docker run -ti -p 5000:5000 --gpus all -e HF_MODEL_ID=nlpconnect/vit-gpt2-image-captioning -e HF_TASK=image-to-text starlette-transformers:gpu
+docker run -ti -p 5000:5000 --gpus all -e HF_MODEL_ID=echarlaix/tiny-random-stable-diffusion-xl -e HF_TASK=text-to-image starlette-transformers:gpu
+docker run -ti -p 5000:5000 --gpus all -e HF_MODEL_ID=stabilityai/stable-diffusion-xl-base-1.0 -e HF_TASK=text-to-image starlette-transformers:gpu
 docker run -ti -p 5000:5000 -e HF_MODEL_DIR=/repository -v $(pwd)/distilbert-base-uncased-emotion:/repository starlette-transformers:cpu
 ```
 

dockerfiles/pytorch/cpu/Dockerfile

@@ -23,7 +23,7 @@ RUN apt-get update \
 # install micromamba
 ENV MAMBA_ROOT_PREFIX=/opt/conda
 ENV PATH=/opt/conda/bin:$PATH
-RUN curl -L https://micromamba.snakepit.net/api/micromamba/linux-64/latest | tar -xj "bin/micromamba" \
+RUN curl -L https://micro.mamba.pm/api/micromamba/linux-64/latest | tar -xj "bin/micromamba" \
     && touch /root/.bashrc \
     && ./bin/micromamba shell init -s bash -p /opt/conda  \
     && grep -v '[ -z "\$PS1" ] && return' /root/.bashrc  > /opt/conda/bashrc

dockerfiles/pytorch/gpu/Dockerfile

@@ -25,7 +25,7 @@ ENV MAMBA_ROOT_PREFIX=/opt/conda
 ENV PATH=/opt/conda/bin:$PATH
 ENV LD_LIBRARY_PATH="/opt/conda/lib:${LD_LIBRARY_PATH}"
 
-RUN curl -L https://micromamba.snakepit.net/api/micromamba/linux-64/latest | tar -xj "bin/micromamba" \
+RUN curl -L https://micro.mamba.pm/api/micromamba/linux-64/latest | tar -xj "bin/micromamba" \
     && touch /root/.bashrc \
     && ./bin/micromamba shell init -s bash -p /opt/conda  \
     && grep -v '[ -z "\$PS1" ] && return' /root/.bashrc  > /opt/conda/bashrc

dockerfiles/pytorch/gpu/environment.yaml

@@ -6,7 +6,7 @@ dependencies:
 - nvidia::cudatoolkit=11.7
 - pytorch::pytorch=1.13.1=py3.9_cuda11.7*
 - pip:
-- transformers[sklearn,sentencepiece,audio,vision]==4.31.0
+  - transformers[sklearn,sentencepiece,audio,vision]==4.31.0
   - sentence_transformers==2.2.2
   - torchvision==0.14.1
   - diffusers==0.19.3

dockerfiles/tensorflow/cpu/Dockerfile

@@ -23,7 +23,7 @@ RUN apt-get update \
 # install micromamba
 ENV MAMBA_ROOT_PREFIX=/opt/conda
 ENV PATH=/opt/conda/bin:$PATH
-RUN curl -L https://micromamba.snakepit.net/api/micromamba/linux-64/latest | tar -xj "bin/micromamba" \
+RUN curl -L https://micro.mamba.pm/api/micromamba/linux-64/latest | tar -xj "bin/micromamba" \
     && touch /root/.bashrc \
     && ./bin/micromamba shell init -s bash -p /opt/conda  \
     && grep -v '[ -z "\$PS1" ] && return' /root/.bashrc  > /opt/conda/bashrc

dockerfiles/tensorflow/gpu/Dockerfile

@@ -26,7 +26,7 @@ ENV MAMBA_ROOT_PREFIX=/opt/conda
 ENV PATH=/opt/conda/bin:$PATH
 ENV LD_LIBRARY_PATH="/opt/conda/lib:${LD_LIBRARY_PATH}"
 
-RUN curl -L https://micromamba.snakepit.net/api/micromamba/linux-64/latest | tar -xj "bin/micromamba" \
+RUN curl -L https://micro.mamba.pm/api/micromamba/linux-64/latest | tar -xj "bin/micromamba" \
     && touch /root/.bashrc \
     && ./bin/micromamba shell init -s bash -p /opt/conda  \
     && grep -v '[ -z "\$PS1" ] && return' /root/.bashrc  > /opt/conda/bashrc

src/huggingface_inference_toolkit/diffusers_utils.py

@@ -43,8 +43,11 @@ def __call__(
             logger.warning("Sending num_images_per_prompt > 1 to pipeline is not supported. Using default value 1.")
 
         # Call pipeline with parameters
-        out = self.pipeline(prompt, num_images_per_prompt=1)
-
+        if self.pipeline.device.type == "cuda":
+            with torch.autocast("cuda"):
+                out = self.pipeline(prompt, num_images_per_prompt=1)
+        else:
+            out = self.pipeline(prompt, num_images_per_prompt=1)
         return out.images[0]
 
 

src/huggingface_inference_toolkit/utils.py

@@ -4,14 +4,13 @@
 from pathlib import Path
 from typing import Optional, Union
 
-from huggingface_hub import login, snapshot_download
+from huggingface_hub import HfApi, login, snapshot_download
 from transformers import WhisperForConditionalGeneration, pipeline
 from transformers.file_utils import is_tf_available, is_torch_available
 from transformers.pipelines import Conversation, Pipeline
 
 from huggingface_inference_toolkit.const import HF_DEFAULT_PIPELINE_NAME, HF_MODULE_NAME
 from huggingface_inference_toolkit.diffusers_utils import (
-    check_supported_pipeline,
     get_diffusers_pipeline,
     is_diffusers_available,
 )
@@ -46,11 +45,12 @@ def is_optimum_available():
     "pt": "pytorch*",
     "flax": "flax*",
     "rust": "rust*",
-    "onnx": "*onnx",
+    "onnx": "*onnx*",
     "safetensors": "*safetensors",
     "coreml": "*mlmodel",
     "tflite": "*tflite",
     "savedmodel": "*tar.gz",
+    "openvino": "*openvino*",
     "ckpt": "*ckpt",
 }
 
@@ -59,18 +59,8 @@ 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 = [
-        "pytorch*",
-        "tf*",
-        "flax*",
-        "rust*",
-        "*onnx",
-        "*safetensors",
-        "*mlmodel",
-        "*tflite",
-        "*tar.gz",
-        "*ckpt",
-    ]
+    ignore_regex_list = list(framework2weight.values())
+
     pattern = framework2weight.get(framework, None)
     if pattern in ignore_regex_list:
         ignore_regex_list.remove(pattern)
@@ -157,6 +147,13 @@ def _load_repository_from_hf(
     if not target_dir.exists():
         target_dir.mkdir(parents=True)
 
+    # check if safetensors weights are available
+    if framework == "pytorch":
+        files = HfApi().model_info(repository_id).siblings
+        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) }")
@@ -259,7 +256,7 @@ def get_pipeline(task: str, model_dir: Path, **kwargs) -> Pipeline:
         "sentence-ranking",
     ]:
         hf_pipeline = get_sentence_transformers_pipeline(task=task, model_dir=model_dir, device=device, **kwargs)
-    elif is_diffusers_available() and check_supported_pipeline(model_dir) and task == "text-to-image":
+    elif is_diffusers_available() and task == "text-to-image":
         hf_pipeline = get_diffusers_pipeline(task=task, model_dir=model_dir, device=device, **kwargs)
     else:
         hf_pipeline = pipeline(task=task, model=model_dir, device=device, **kwargs)