Update handler.py

handler.py

@@ -1,35 +1,47 @@
-# handler.py
-import torch
-from torchvision import transforms
-from PIL import Image
-from huggingface_hub import hf_hub_download
-import io
-import base64
-import json
-
-class EndpointHandler:
-    def __init__(self, path=""):
-        model_path = hf_hub_download(repo_id="whitney0507/unet-model", filename="UNet_Model.pth")
-        self.model = torch.load(model_path, map_location="cpu")
-        self.model.eval()
-        self.transform = transforms.Compose([
-            transforms.Resize((256, 256)),
-            transforms.ToTensor()
-        ])
-
-    def __call__(self, data):
-        image_bytes = base64.b64decode(data["inputs"])
-        image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
-        input_tensor = self.transform(image).unsqueeze(0)
-
-        with torch.no_grad():
-            output = self.model(input_tensor)
-            pred = output.argmax(dim=1).squeeze().byte().cpu().numpy()
-
-        # Convert output back to base64 image
-        pred_img = Image.fromarray(pred * 255)
-        buffer = io.BytesIO()
-        pred_img.save(buffer, format="PNG")
-        base64_output = base64.b64encode(buffer.getvalue()).decode("utf-8")
-
-        return { "prediction": base64_output }
+# handler.py
+# Save only the weights
+torch.save(model.state_dict(), "UNet_Model.pth")
+import torch
+import torch.nn as nn
+from torchvision import transforms
+from PIL import Image
+from huggingface_hub import hf_hub_download
+import io
+import base64
+
+# Define your UNet class here (shortened version for example)
+class UNet(nn.Module):
+    def __init__(self):  # Add your actual init params
+        super(UNet, self).__init__()
+        # Define layers...
+
+    def forward(self, x):
+        # Implement forward pass
+        return x
+
+class EndpointHandler:
+    def __init__(self, path=""):
+        model_path = hf_hub_download(repo_id="whitney0507/unet-model", filename="UNet_Model.pth")
+        self.model = UNet()  # Instantiate model
+        self.model.load_state_dict(torch.load(model_path, map_location="cpu"))
+        self.model.eval()
+        self.transform = transforms.Compose([
+            transforms.Resize((256, 256)),
+            transforms.ToTensor()
+        ])
+
+    def __call__(self, data):
+        image_bytes = base64.b64decode(data["inputs"])
+        image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
+        input_tensor = self.transform(image).unsqueeze(0)
+
+        with torch.no_grad():
+            output = self.model(input_tensor)
+            pred = output.argmax(dim=1).squeeze().byte().cpu().numpy()
+
+        # Convert to base64
+        output_img = Image.fromarray(pred * 255)
+        buffer = io.BytesIO()
+        output_img.save(buffer, format="PNG")
+        return {"prediction": base64.b64encode(buffer.getvalue()).decode("utf-8")}
+