Update handler.py

handler.py

@@ -2,7 +2,6 @@ import torch
 from transformers import CLIPProcessor, CLIPModel
 from PIL import Image
 import requests
-import base64
 import io
 
 class EndpointHandler:
@@ -11,61 +10,35 @@ class EndpointHandler:
         self.model = CLIPModel.from_pretrained("dazpye/clip-image")
         self.processor = CLIPProcessor.from_pretrained("dazpye/clip-image")
 
-    def _load_image(self, image_data):
-        """Fetches an image from a URL or decodes a base64 image."""
+    def _load_image(self, image_url):
+        """Simple image loader for URL images."""
         try:
-            if isinstance(image_data, str):
-                if image_data.startswith("http"):  
-                    # Fetch image from URL with User-Agent to bypass restrictions
-                    print(f"🌐 Fetching image from: {image_data}")
-                    headers = {
-                        "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36"
-                    }
-                    response = requests.get(image_data, headers=headers, timeout=5)
-                    print(f"✅ HTTP Status Code: {response.status_code}")
-
-                    if response.status_code == 200:
-                        image_bytes = io.BytesIO(response.content)
-                        return Image.open(image_bytes).convert("RGB")
-                    else:
-                        print(f"❌ Failed to fetch image: HTTP {response.status_code}")
-                
-                else:
-                    # Handle base64-encoded image
-                    print("📸 Decoding base64 image...")
-                    return Image.open(io.BytesIO(base64.b64decode(image_data))).convert("RGB")
-
+            print(f"🌐 Fetching image: {image_url}")
+            response = requests.get(image_url, timeout=5)
+            response.raise_for_status()  # Raise error if status is not 200
+            return Image.open(io.BytesIO(response.content)).convert("RGB")
         except Exception as e:
-            print(f"⚠️ Exception in image loading: {e}")
-        
+            print(f"❌ Image loading failed: {e}")
         return None  # Return None if image loading fails
 
     def __call__(self, data):
-        """Main inference function Hugging Face expects."""
-        print("📥 Processing input...")
+        """Processes input and runs inference."""
+        print("📥 Received input...")
 
-        if "inputs" in data:
-            data = data["inputs"]
+        text = data.get("inputs", {}).get("text", ["default text"])
+        image_urls = data.get("inputs", {}).get("images", [])
 
-        text = data.get("text", ["default caption"])  # Default text
-        images = data.get("images", [])  # List of images
+        images = [self._load_image(url) for url in image_urls if url]
+        images = [img for img in images if img]  # Remove failed images
 
-        # Convert image URLs or base64 strings to PIL images
-        pil_images = [self._load_image(img) for img in images if img]
-        pil_images = [img for img in pil_images if img]  # Remove None values
+        if not images:
+            print("❌ No valid images provided.")
+            return {"error": "No valid images provided."}
 
-        if not pil_images:  
-            print("❌ No valid images provided. Check URLs or base64 encoding.")
-            return {"error": "❌ No valid images provided. Check URLs or base64 encoding."}
-
-        inputs = self.processor(text=text, images=pil_images, return_tensors="pt")
+        inputs = self.processor(text=text, images=images, return_tensors="pt")
 
         print("🖥️ Running inference...")
         with torch.no_grad():
             outputs = self.model(**inputs)
 
-        logits_per_image = outputs.logits_per_image
-        probabilities = logits_per_image.softmax(dim=1)
-
-        print("✅ Inference complete!")
-        return {"predictions": probabilities.tolist()}
+        return {"predictions": outputs.logits_per_image.softmax(dim=1).tolist()}