Create handler.py

handler.py

@@ -0,0 +1,42 @@
+import base64
+from io import BytesIO
+from PIL import Image
+import torch
+from transformers import CLIPProcessor, CLIPModel
+from typing import Dict, Any
+
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+class EndpointHandler():
+    def __init__(self, path=""):
+        self.processor = CLIPProcessor.from_pretrained("openai/openai/clip-vit-large-patch14")
+        self.model = CLIPModel.from_pretrained("openai/openai/clip-vit-large-patch14").to(device)
+        self.model.eval()
+
+    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
+        input_data = data.get("inputs", {})
+        encoded_images = input_data.get("images")
+        texts = input_data.get("texts", [])
+
+        if not encoded_images or not texts:
+            return {"error": "Both images and texts must be provided"} 
+
+        try:
+            images = [Image.open(BytesIO(base64.b64decode(img))).convert("RGB") for img in encoded_images]
+            inputs = self.processor(text=texts, images=images, return_tensors="pt", padding=True)
+
+            # Move tensors to the same device as model
+            inputs = {k: v.to(device) for k, v in inputs.items()}
+
+            with torch.no_grad():
+                outputs = self.model(**inputs)
+                logits_per_image = outputs.logits_per_image  # this is the image-text similarity score
+                logits_per_text = outputs.logits_per_text  # this is the text-image similarity score
+
+            return {
+                "logits_per_image": logits_per_image.cpu().numpy().tolist(),
+                "logits_per_text": logits_per_text.cpu().numpy().tolist()
+            }
+        except Exception as e:
+            print(f"Error during processing: {str(e)}")
+            return {"error": str(e)}