Added custom handler

handler.py

@@ -1,81 +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()}
+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

@@ -1,48 +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"]
-# )
+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"]
+# )