custom handler

handler.py

@@ -0,0 +1,53 @@
+from typing import Dict, List, Any
+from PIL import Image
+from transformers import AutoProcessor, AutoModelForVision2Seq
+import base64
+from io import BytesIO
+
+
+class EndpointHandler():
+    def __init__(self, path=""):
+        # Preload all the elements you are going to need at inference.
+        # pseudo:
+        # self.model= load_model(path)
+        self.model = AutoModelForVision2Seq.from_pretrained("microsoft/kosmos-2-patch14-224").to("cuda")
+        self.processor = AutoProcessor.from_pretrained("microsoft/kosmos-2-patch14-224")
+
+        # prompt = "<grounding>An image of"
+
+    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
+        
+        prompt = data.pop("prompt")
+        image_base64 = data.pop("image_base64")
+
+        image_data = base64.b64decode(image_base64)
+        image = Image.open(BytesIO(image_data))
+
+        inputs = self.processor(text=prompt, images=image, return_tensors="pt").to("cuda")
+
+        generated_ids = self.model.generate(
+            pixel_values=inputs["pixel_values"],
+            input_ids=inputs["input_ids"],
+            attention_mask=inputs["attention_mask"],
+            image_embeds=None,
+            image_embeds_position_mask=inputs["image_embeds_position_mask"],
+            use_cache=True,
+            max_new_tokens=128,
+        )
+        generated_text = self.processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
+
+        # Specify `cleanup_and_extract=False` in order to see the raw model generation.
+        processed_text = self.processor.post_process_generation(generated_text, cleanup_and_extract=False)
+
+        # print(processed_text)
+        # `<grounding> An image of<phrase> a snowman</phrase><object><patch_index_0044><patch_index_0863></object> warming himself by<phrase> a fire</phrase><object><patch_index_0005><patch_index_0911></object>.`
+
+        # By default, the generated  text is cleanup and the entities are extracted.
+        processed_text, entities = self.processor.post_process_generation(generated_text)
+
+        # print(processed_text)
+        # `An image of a snowman warming himself by a fire.`
+        return [{"processed_text": processed_text}]
+
+        # print(entities)
+        # `[('a snowman', (12, 21), [(0.390625, 0.046875, 0.984375, 0.828125)]), ('a fire', (41, 47), [(0.171875, 0.015625, 0.484375, 0.890625)])]`

requirements.txt

@@ -0,0 +1,2 @@
+base64
+pillow

test_handler.py

@@ -0,0 +1,40 @@
+from handler import EndpointHandler
+from PIL import Image
+import requests
+import base64
+from io import BytesIO
+import time
+
+# init handler
+my_handler = EndpointHandler(path=".")
+
+
+
+# prompt = "<grounding>An image of"
+prompt = "<grounding>Describe this image in detail"
+
+url = "https://huggingface.co/microsoft/kosmos-2-patch14-224/resolve/main/snowman.png"
+image = Image.open(requests.get(url, stream=True).raw)
+
+# The original Kosmos-2 demo saves the image first then reload it. For some images, this will give slightly different image input and change the generation outputs.
+image.save("new_image.jpg")
+image = Image.open("img1.jpg")
+buffered = BytesIO()
+image.save(buffered, format='JPEG')
+image_base64 = base64.b64encode(buffered.getvalue()).decode('utf-8')
+
+# prepare sample payload
+non_holiday_payload = {"prompt": prompt, "image_base64": image_base64}
+# holiday_payload = {"inputs": "Today is a though day", "date": "2022-07-04"}
+init_t = time.time()
+# test the handler
+non_holiday_pred=my_handler(non_holiday_payload)
+# holiday_payload=my_handler(holiday_payload)
+
+# show results
+print("image_description", non_holiday_pred)
+print(time.time() - init_t)
+# print("holiday_payload", holiday_payload)
+
+# non_holiday_pred [{'label': 'joy', 'score': 0.9985942244529724}]
+# holiday_payload [{'label': 'happy', 'score': 1}]