Update handler.py

handler.py

@@ -1,8 +1,7 @@
 import torch
-from transformers import pipeline
-from PIL import Image
-import requests
 import os
+from huggingface_hub import PyTorchModelHubMixin
+
 
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') # Set device
 print('device >>> ', device)
@@ -10,19 +9,35 @@ print('device >>> ', device)
 
 class EndpointHandler():
 
+    class MyModel(
+        nn.Module,
+        PyTorchModelHubMixin,
+    ):
+    def __init__(self):
+        super().__init__()
+        self.model = model
+
+
+    def forward(self, x):
+        x = self.model(x)
+        return x
+
+    net = MyModel()
+
     def __init__(self, path=""):
-        #pipe = pipeline("image-segmentation", "nvidia/segformer-b1-finetuned-cityscapes-1024-1024") #semantic_segmentation
-        # pipe = pipeline("image-segmentation", "facebook/mask2former-swin-large-cityscapes-instance") #instance_segmentation
-        pipe = pipeline("image-segmentation", "facebook/mask2former-swin-large-cityscapes-panoptic") #panoptic_segmentation
-        self.model = pipe # No need to move model to device for this algo
+        model = MyModel.from_pretrained("damiano216/pay-boo-2")
+        self.model = model 
+        #POTENTIALLY WILL NEED TO MOVE THE MODEL TO DEVICE HERE
 
 
     def __call__(self, data):
-        url = data["image_url"]
-        image = Image.open(requests.get(url, stream=True).raw)
-        print('image >>> ', image)
 
-        results = self.model(image)
-        print('results >>> ', results)
+        new_data_tensor = data['chargeData']
+        # 3. Make predictions
+        with torch.no_grad():
+            predictions = self.model(new_data_tensor)
+
+        # 4. Interpret predictions
+        print(f"predictions >>> : {predictions[0][0]}")
 
-        return results
+        return predictions