Update handler.py

handler.py

@@ -1,32 +1,44 @@
 import torch
 from transformers import CLIPProcessor, CLIPModel
+from PIL import Image
+import base64
+import io
 
 class EndpointHandler:
-    def __init__(self):
-        # Load model and processor
+    def __init__(self, model_dir=None):  # AWS expects model_dir
+        print("Loading model...")
         self.model = CLIPModel.from_pretrained("dazpye/clip-image")
         self.processor = CLIPProcessor.from_pretrained("dazpye/clip-image")
 
-    def preprocess(self, inputs):
-        # Process input data
-        text = inputs.get("text", [])
-        images = inputs.get("images", [])
-        return self.processor(text=text, images=images, return_tensors="pt")
+    def _load_image(self, image_data):
+        """Handles both URL-based and base64 image inputs."""
+        if isinstance(image_data, str):  
+            if image_data.startswith("http"):  
+                return Image.open(requests.get(image_data, stream=True).raw)
+            else:  # Assume base64-encoded image
+                return Image.open(io.BytesIO(base64.b64decode(image_data)))
+        return None  # Invalid image format
 
-    def inference(self, inputs):
-        # Run inference
+    def __call__(self, data):
+        """Main inference function Hugging Face expects."""
+        print("Processing input...")
+        
+        text = data.get("text", ["default caption"])  # Default text
+        images = data.get("images", [])  # List of images
+        
+        # Convert image URLs or base64 strings to PIL images
+        pil_images = [self._load_image(img) for img in images if img]
+
+        if not pil_images:  
+            return {"error": "No valid images provided."}
+        
+        inputs = self.processor(text=text, images=pil_images, return_tensors="pt")
+
+        print("Running inference...")
         with torch.no_grad():
             outputs = self.model(**inputs)
-        return outputs.logits_per_image.tolist()
-
-    def postprocess(self, inference_output):
-        # Convert output to readable format
-        return {"predictions": inference_output}
 
-handler = EndpointHandler()
+        logits_per_image = outputs.logits_per_image
+        probabilities = logits_per_image.softmax(dim=1)
 
-def handle(request):
-    inputs = request if isinstance(request, dict) else request.json()
-    processed_inputs = handler.preprocess(inputs)
-    predictions = handler.inference(processed_inputs)
-    return handler.postprocess(predictions)
+        return {"predictions": probabilities.tolist()}