Update handler.py

handler.py

@@ -1,79 +1,15 @@
-import torch
 from transformers import AutoFeatureExtractor, AutoModelForImageClassification
-from PIL import Image
-import json
-import base64
-from io import BytesIO
 
-# Load the model and feature extractor when the handler is initialized
 class VisionModelHandler:
     def __init__(self, model_name_or_path="https://huggingface.co/VishalD1234/Florence-metere1"):
         self.model_name_or_path = model_name_or_path
         self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         
-        # Load model and feature extractor
-        self.model = AutoModelForImageClassification.from_pretrained(self.model_name_or_path)
-        self.feature_extractor = AutoFeatureExtractor.from_pretrained(self.model_name_or_path)
+        # Add the trust_remote_code parameter
+        self.model = AutoModelForImageClassification.from_pretrained(self.model_name_or_path, trust_remote_code=True)
+        self.feature_extractor = AutoFeatureExtractor.from_pretrained(self.model_name_or_path, trust_remote_code=True)
         
-        # Move the model to the appropriate device (GPU/CPU)
         self.model.to(self.device)
-        self.model.eval()  # Set model to evaluation mode
-    
-    def preprocess_image(self, image_data):
-        """
-        Preprocess the image for the model. Convert it from base64 and apply the necessary transformations.
-        """
-        image = Image.open(BytesIO(base64.b64decode(image_data)))
-        inputs = self.feature_extractor(images=image, return_tensors="pt").to(self.device)
-        return inputs
-    
-    def predict(self, inputs):
-        """
-        Perform inference and return the model's predictions.
-        """
-        with torch.no_grad():
-            outputs = self.model(**inputs)
-            logits = outputs.logits
-            predicted_class_idx = logits.argmax(-1).item()  # Get the index of the highest score
-        return predicted_class_idx
-    
-    def handle(self, event, context):
-        """
-        Entry point for the inference request. This will be called by the inference endpoint.
-        
-        Args:
-            event: This will contain the input data, usually in the form of a JSON with an image in base64.
-            context: Optional, can contain metadata about the request (not used here).
-        
-        Returns:
-            A JSON response with the prediction result.
-        """
-        # Extract image data from the request body
-        body = json.loads(event.get("body", "{}"))
-        image_data = body.get("image_base64", None)
-        
-        if image_data is None:
-            return {
-                "statusCode": 400,
-                "body": json.dumps({"error": "No image data found in the request"})
-            }
-        
-        # Preprocess the image and make predictions
-        inputs = self.preprocess_image(image_data)
-        prediction = self.predict(inputs)
-        
-        # You can add more details to this mapping depending on your use case
-        response = {
-            "statusCode": 200,
-            "body": json.dumps({"prediction": prediction})
-        }
-        
-        return response
-
-
-# Instantiate the handler
-vision_model_handler = VisionModelHandler()
+        self.model.eval()
 
-# If running in an environment like AWS Lambda or Sagemaker, ensure this is exposed
-def lambda_handler(event, context):
-    return vision_model_handler.handle(event, context)
+    # Rest of the code stays the same...