Create handler.py

handler.py

@@ -0,0 +1,59 @@
+import torch
+from PIL import Image
+import base64
+from io import BytesIO
+from transformers import AutoModel, AutoTokenizer
+
+class EndpointHandler:
+    def __init__(self, path="/repository"):
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        
+        # Load the model
+        self.model = AutoModel.from_pretrained(
+            path,
+            trust_remote_code=True,
+            attn_implementation='sdpa',
+            torch_dtype=torch.bfloat16 if self.device.type == "cuda" else torch.float32,
+        ).to(self.device)
+        self.model.eval()
+        
+        # Load the tokenizer
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            path,
+            trust_remote_code=True,
+        )
+
+    def __call__(self, data):
+        # Extract image and text from the input data
+        image_data = data.get("inputs", {}).get("image", "")
+        text_prompt = data.get("inputs", {}).get("text", "")
+
+        if not image_data or not text_prompt:
+            return {"error": "Both 'image' and 'text' must be provided in the input data."}
+
+        # Process the image data
+        try:
+            image_bytes = base64.b64decode(image_data)
+            image = Image.open(BytesIO(image_bytes)).convert("RGB")
+        except Exception as e:
+            return {"error": f"Failed to process image data: {e}"}
+
+        # Prepare the messages for the model
+        msgs = [{'role': 'user', 'content': [image, text_prompt]}]
+
+        # Generate output
+        with torch.no_grad():
+            res = self.model.chat(
+                image=None,
+                msgs=msgs,
+                tokenizer=self.tokenizer,
+                sampling=True,
+                temperature=0.7,
+                top_p=0.95,
+                max_length=2000,
+            )
+        
+        # The result is the generated text
+        output_text = res
+
+        return {"generated_text": output_text}