Add config.json with model_type

config.json

@@ -4,7 +4,7 @@
         "U2NET"
     ],
     "task": "image-segmentation",
-    "pipeline_class": "U2NetPipeline",
+    "pipeline_class": "u2net_pipeline.U2NetPipeline",
     "model_file": "u2net.py",
     "weights_file": "u2net.pth"
 }

u2net_pipeline.py

@@ -10,10 +10,15 @@ from transformers import Pipeline
 class U2NetPipeline(Pipeline):
     def __init__(self, model, **kwargs):
         super().__init__(model=model, **kwargs)
-        self.model = U2NET(3, 1)
-        self.model.load_state_dict(torch.load(model, map_location="cpu"))
+        self.model = model
         self.model.eval()
 
+    @classmethod
+    def from_pretrained(cls, model_path, **kwargs):
+        model = U2NET(3, 1)
+        model.load_state_dict(torch.load(f"{model_path}/u2net.pth", map_location="cpu"))
+        return cls(model, **kwargs)
+
     def _sanitize_parameters(self, **kwargs):
         return {}, {}, {}
 
@@ -50,5 +55,6 @@ class U2NetPipeline(Pipeline):
         result = (result - mi) / (ma - mi)
         return (result * 255).astype(np.uint8)
 
-def load_model():
-    return U2NetPipeline("u2net.pth")
+# Remove or comment out this function as it's no longer needed
+# def load_model():
+#     return U2NetPipeline("u2net.pth")