added custom handler and test script

handler.py

@@ -0,0 +1,81 @@
+from typing import Dict, List, Any
+from transformers import pipeline
+from PIL import Image
+import requests
+import os
+from io import BytesIO
+from transformers import CLIPSegProcessor, CLIPSegForImageSegmentation
+from diffusers import DiffusionPipeline
+import torch
+from torch import autocast
+import base64
+
+
+auth_token = "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"
+
+
+class EndpointHandler():
+    def __init__(self, path=""):
+        self.processor = CLIPSegProcessor.from_pretrained("./clipseg-rd64-refined")
+        self.model = CLIPSegForImageSegmentation.from_pretrained("./clipseg-rd64-refined")
+
+        self.pipe = DiffusionPipeline.from_pretrained(
+            "./",
+            custom_pipeline="text_inpainting",
+            segmentation_model=self.model,
+            segmentation_processor=self.processor,
+            revision="fp16",
+            torch_dtype=torch.float16,
+            use_auth_token=auth_token,
+        )
+
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.pipe = self.pipe.to(self.device)
+
+    def pad_image(self, image):
+        w, h = image.size
+        if w == h:
+            return image
+        elif w > h:
+            new_image = Image.new(image.mode, (w, w), (0, 0, 0))
+            new_image.paste(image, (0, (w - h) // 2))
+            return new_image
+        else:
+            new_image = Image.new(image.mode, (h, h), (0, 0, 0))
+            new_image.paste(image, ((h - w) // 2, 0))
+            return new_image
+
+
+    def process_image(self, image, text, prompt):
+        image = self.pad_image(image)
+        image = image.resize((512, 512))
+        with autocast(self.device):
+            inpainted_image = self.pipe(image=image, text=text, prompt=prompt).images[0]
+        return inpainted_image
+
+    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
+        """
+       data args:
+            inputs (:obj: `str`)
+            date (:obj: `str`)
+      Return:
+            A :obj:`list` | `dict`: will be serialized and returned
+        """
+        # get inputs
+        inputs = data.pop("inputs", data)
+
+        # decode base64 image to PIL
+        image = Image.open(BytesIO(base64.b64decode(inputs['image'])))
+        class_text = inputs['class_text']
+        prompt = inputs['prompt']
+        # run inference pipeline
+        with autocast(self.device):
+            image = self.process_image(image, class_text, prompt)
+            
+        # encode image as base 64
+        buffered = BytesIO()
+        image.save(buffered, format="JPEG")
+        img_str = base64.b64encode(buffered.getvalue())
+
+        # postprocess the prediction
+        return {"image": img_str.decode()}

test_handler.py

@@ -0,0 +1,48 @@
+from handler import EndpointHandler
+import json
+from typing import List
+import requests as r
+import base64
+import requests as r
+import base64
+from PIL import Image
+from io import BytesIO
+
+ENDPOINT_URL = ""
+HF_TOKEN = ""
+def decode_base64_image(image_string):
+  base64_image = base64.b64decode(image_string)
+  buffer = BytesIO(base64_image)
+  return  Image.open(buffer)
+
+# init handler
+my_handler = EndpointHandler(path=".")
+
+# prepare sample payload
+path_to_image = "test_images/lal.jpg"
+with open(path_to_image, "rb") as i:
+    b64 = base64.b64encode(i.read())
+
+payload = {"inputs": {"image": b64.decode("utf-8"), "class_text": "shirt", "prompt": "wedding shirt"}}
+
+# test the handler
+results=my_handler(payload)
+
+# show results
+# print("non_holiday_pred", non_holiday_pred)
+# print("holiday_payload", holiday_payload)
+decode_base64_image(results["image"]).save("test_results.jpg")
+# def predict(path_to_image: str = None, candiates: List[str] = None):
+#     with open(path_to_image, "rb") as i:
+#         b64 = base64.b64encode(i.read())
+
+#     payload = {"inputs": {"image": b64.decode("utf-8"), "candiates": candiates}}
+#     response = r.post(
+#         ENDPOINT_URL, headers={"Authorization": f"Bearer {HF_TOKEN}"}, json=payload
+#     )
+#     return response.json()
+
+
+# prediction = predict(
+#     path_to_image="palace.jpg", candiates=["sea", "palace", "car", "ship"]
+# )