add pipeline

MangaPipe.py

@@ -0,0 +1,22 @@
+from transformers import Pipeline , AutoModel, AutoImageProcessor, pipeline
+from PIL import Image
+from loadimg import load_img
+import torch
+REPO_NAME = "p1atdev/MangaLineExtraction-hf"
+
+class MangaLineExtractionPipe(Pipeline):
+  def __init__(self,**kwargs):
+    self.processor = AutoImageProcessor.from_pretrained("p1atdev/MangaLineExtraction-hf", trust_remote_code=True)
+    Pipeline.__init__(self,**kwargs)
+    self.model.eval()
+  def _sanitize_parameters(kwargs):
+    return {}, {} , {}
+  def preprocess(self,image):
+    # takes any image type and returns a pil image
+    image = load_img(image, output_type = "pil").convert("RGB")
+    processed =  self.processor(image, return_tensors="pt")
+    return processed.pixel_values
+  def _forward(self,inputs):
+    return self.model(inputs)
+  def postprocess(self,outputs):
+    return Image.fromarray(outputs.pixel_values[0].numpy().astype("uint8"), mode="L")

config.json

@@ -6,6 +6,16 @@
     "AutoConfig": "configuration_mle.MLEConfig",
     "AutoModel": "modeling_mle.MLEForAnimeLineExtraction"
   },
+  "custom_pipelines": {
+    "image-to-image": {
+      "impl": "MangaPipe.MangaLineExtractionPipe",
+      "pt": [
+        "AutoModel"
+      ],
+      "tf": [],
+      "type": "image"
+    }
+  },
   "batch_norm_eps": 0.001,
   "block_kernel_size": 3,
   "block_patch_size": 24,