Updated model to use PyTorch instead of ONNX

pipeline.py

@@ -16,7 +16,8 @@ class Pipeline:
         )
         self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
         self.is_torch_script = self.device.type == 'cpu'
-        self.model = self._load_model(model_name)
+        self.checkpoint = f"models/{model_name}.pt"
+        self.model = self._load_model()
 
         self._log_device_info()
 
@@ -33,7 +34,7 @@ class Pipeline:
 
         return np.squeeze(output, axis=0).squeeze()
 
-    def _load_pytorch_model(self, checkpoint):
+    def _load_pytorch_model(self):
         model = SwinMattingModel({
             "encoder": {
                 "model_name": "microsoft/swin-small-patch4-window7-224"
@@ -43,23 +44,23 @@ class Pipeline:
                 "refine_channels": 16
             }
         })
-        self._load_checkpoint(model, checkpoint)
+        self._load_checkpoint(model)
 
         model.to(self.device)
         model.eval()
 
         return model
 
-    def _load_model(self, model_name):
-        model = self._load_pytorch_model(checkpoint=f"models/{model_name}.pt")
+    def _load_model(self):
+        model = self._load_pytorch_model()
 
         model.to(self.device)
         model.eval()
 
         return model
 
-    def _load_checkpoint(self, model, checkpoint_path):
-        checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=True)
+    def _load_checkpoint(self, model):
+        checkpoint = torch.load(self.checkpoint, map_location="cpu", weights_only=True)
 
         missing_keys, unexpected_keys = model.load_state_dict(checkpoint)
         if missing_keys: