Update handler.py

handler.py

@@ -4,12 +4,13 @@ import torch
 import os
 import torchvision.transforms as transforms
 
-class AIImageHandler:
-    def __init__(self):
+class EndpointHandler:
+    def __init__(self, model_dir):
         self.model = None
         self.extractor = None
         self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         self.labels = None
+        self.initialize(model_dir)
 
     def initialize(self, model_dir):
         config = AutoConfig.from_pretrained(model_dir)
@@ -42,4 +43,4 @@ class AIImageHandler:
         image_data = data.get("inputs") or data.get("image")
         if isinstance(image_data, Image.Image):
             return self.predict({"image": image_data})
-        return {"error": "Invalid input"}
+        return {"error": "Invalid input"}