adding image-to-text pipeline

Author: philschmid

README.md

@@ -39,6 +39,7 @@ 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 -e HF_MODEL_DIR=/repository -v $(pwd)/distilbert-base-uncased-emotion:/repository starlette-transformers:cpu
 ```
 

dockerfiles/pytorch/gpu/Dockerfile

@@ -1,4 +1,4 @@
-FROM nvidia/cuda:11.7.0-devel-ubuntu22.04
+FROM nvidia/cuda:11.7.1-devel-ubuntu22.04
 
 LABEL maintainer="Hugging Face"
 

dockerfiles/pytorch/gpu/environment.yaml

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

src/huggingface_inference_toolkit/handler.py

@@ -25,7 +25,6 @@ def __call__(self, data):
         """
         inputs = data.pop("inputs", data)
         parameters = data.pop("parameters", None)
-
         # pass inputs with all kwargs in data
         if parameters is not None:
             prediction = self.pipeline(inputs, **parameters)

src/huggingface_inference_toolkit/utils.py

@@ -243,6 +243,8 @@ def get_pipeline(task: str, model_dir: Path, **kwargs) -> Pipeline:
         "zero-shot-image-classification",
     }:
         kwargs["feature_extractor"] = model_dir
+    elif task in {"image-to-text"}:
+        pass
     else:
         kwargs["tokenizer"] = model_dir
 
@@ -278,3 +280,15 @@ def get_pipeline(task: str, model_dir: Path, **kwargs) -> Pipeline:
             (rank + 1, token) for rank, token in enumerate(hf_pipeline.tokenizer.prefix_tokens[1:])
         ]
     return hf_pipeline
+
+
+def convert_params_to_int_or_bool(params):
+    """Converts query params to int or bool if possible"""
+    for k, v in params.items():
+        if v.isnumeric():
+            params[k] = int(v)
+        if v == 'false':
+            params[k] = False
+        if v == 'true':
+            params[k] = True
+    return params

src/huggingface_inference_toolkit/webservice_starlette.py

@@ -19,7 +19,7 @@
 from huggingface_inference_toolkit.handler import get_inference_handler_either_custom_or_default_handler
 from huggingface_inference_toolkit.serialization.base import ContentType
 from huggingface_inference_toolkit.serialization.json_utils import Jsoner
-from huggingface_inference_toolkit.utils import _load_repository_from_hf
+from huggingface_inference_toolkit.utils import _load_repository_from_hf, convert_params_to_int_or_bool
 
 
 def config_logging(level=logging.INFO):
@@ -64,8 +64,6 @@ async def health(request):
 
 async def predict(request):
     try:
-        # tracks request time
-        start_time = perf_counter()
         # extracts content from request
         content_type = request.headers.get("content-Type", None)
         # try to deserialize payload
@@ -74,13 +72,16 @@ async def predict(request):
         if "inputs" not in deserialized_body:
             raise ValueError(f"Body needs to provide a inputs key, recieved: {orjson.dumps(deserialized_body)}")
 
+        # check for query parameter and add them to the body
+        if request.query_params and "parameters" not in deserialized_body:
+            deserialized_body["parameters"] = convert_params_to_int_or_bool(dict(request.query_params))
+        print(deserialized_body)
+        
+        # tracks request time
+        start_time = perf_counter()
         # run async not blocking call
         pred = await async_handler_call(inference_handler, deserialized_body)
-        # run sync blocking call -> slighty faster for < 200ms prediction time
-        # pred = inference_handler(deserialized_body)
-
         # log request time
-        # TODO: repalce with middleware
         logger.info(f"POST {request.url.path} | Duration: {(perf_counter()-start_time) *1000:.2f} ms")
 
         # response extracts content from request