Add custom `handler.py`

handler.py

@@ -0,0 +1,59 @@
+from typing import Any, Dict, List
+
+import requests
+import torch
+
+from transformers import AutoProcessor, PaliGemmaForConditionalGeneration
+from PIL import Image
+
+
+class EndpointHandler:
+    def __init__(
+        self,
+        model_dir: str = "/opt/huggingface/model",
+        **kwargs: Any,
+    ) -> None:
+        self.model = PaliGemmaForConditionalGeneration.from_pretrained(
+            "google/paligemma-3b-mix-448",
+            revision="bfloat16",
+            torch_dtype=torch.bfloat16,
+            low_cpu_mem_usage=True,
+            device_map="auto",
+        ).eval()
+
+        self.processor = AutoProcessor.from_pretrained("google/paligemma-3b-mix-448")
+
+    def __call__(self, data: Dict[str, Any]) -> Dict[str, List[Any]]:
+        if "instances" not in data:
+            raise ValueError(
+                "The request body must contain a key `instances` with a list of instances."
+            )
+
+        predictions = []
+        for input in data["instances"]:
+            if any(key not in input for key in {"prompt", "image_url"}):
+                raise ValueError(
+                    "The request body for each instance should contain both the `prompt` and the `image_url` key with a valid image URL."
+                )
+
+            try:
+                image = Image.open(requests.get(input["image_url"], stream=True).raw)  # type: ignore
+            except Exception as e:
+                raise ValueError(
+                    f"The provided image URL ({input['image_url']}) cannot be downloaded (with exception {e}), make sure it's public and accessible."
+                )
+
+            inputs = self.processor(
+                text=input["prompt"], images=image, return_tensors="pt"
+            ).to(self.model.device)
+            input_len = inputs["input_ids"].shape[-1]
+
+            with torch.inference_mode():
+                generation_kwargs = data.get(
+                    "generation_kwargs", {"max_new_tokens": 100, "do_sample": False}
+                )
+                generation = self.model.generate(**inputs, **generation_kwargs)
+                generation = generation[0][input_len:]
+                response = self.processor.decode(generation, skip_special_tokens=True)
+            predictions.append(response)
+        return {"predictions": predictions}