Update handler.py

handler.py

@@ -1,8 +1,9 @@
+import base64
 from io import BytesIO
 from typing import Any, List, Dict
 
 from PIL import Image
-from transformers import AutoModel, AutoTokenizer
+from transformers import AutoTokenizer, AutoModel
 
 
 class EndpointHandler():
@@ -16,18 +17,26 @@ class EndpointHandler():
         self.tokenizer = tokenizer
 
     def __call__(self, data: Any) -> List[List[Dict[str, float]]]:
-        image_bytes = data.pop("image_bytes", None)
-        question = data.pop("question", None)
-        image = Image.open(BytesIO(image_bytes))
+        inputs = data.pop("inputs", data)
 
-        msgs = [{'role': 'user', 'content': question}]
+        image = inputs.pop("image", None)  #  base64 image as bytes
+        question = inputs.pop("question", None)
+        msgs = inputs.pop("msgs", None)
+
+
+        parameters = data.pop("parameters", {})
+
+        image = Image.open(BytesIO(base64.b64decode(image)))
+
+        if not msgs:
+            msgs = [{'role': 'user', 'content': question}]
 
         res = self.model.chat(
             image=image,
             msgs=msgs,
             tokenizer=self.tokenizer,
             sampling=True,  # if sampling=False, beam_search will be used by default
-            temperature=0.7,
+            temperature=parameters.get('temperature', 0.7),
             # system_prompt='' # pass system_prompt if needed
         )