Create handler.py

handler.py

@@ -0,0 +1,34 @@
+import base64
+import requests
+from PIL import Image
+from typing import Dict, List, Any, Union
+import torch
+from io import BytesIO
+from transformers import BlipProcessor, BlipForConditionalGeneration, BitsAndBytesConfig
+
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+class EndpointHandler():
+    def __init__(self, model_dir="Salesforce/blip-image-captioning-large"):
+        self.model = BlipForConditionalGeneration.from_pretrained(model_dir).to(device).eval()
+        self.processor = BlipProcessor.from_pretrained(model_dir)
+
+    def __call__(self, data):
+        img_url = data.get('img_url')
+        text_prompt = data.get('text', None)
+        
+        raw_image = Image.open(requests.get(img_url, stream=True).raw).convert('RGB')
+        
+        if text_prompt:
+            inputs = self.processor(raw_image, text_prompt, return_tensors="pt").to(device)
+        else:
+            inputs = self.processor(raw_image, return_tensors="pt").to(device)
+        
+        with torch.no_grad():
+            generated_ids = self.model.generate(
+                **inputs, 
+                max_new_tokens=150
+            )
+        generated_text = self.processor.decode(generated_ids[0], skip_special_tokens=True)
+
+        return {"responses": generated_text}