handler.py

handler.py

@@ -1,28 +1,36 @@
-from typing import Dict, List, Any
-from transformers import pipeline
 
-class PreTrainedPipeline():
+from typing import  Dict, List, Any
+import torch
+from torch import autocast
+from diffusers import StableDiffusionPipeline
+import base64
+from io import BytesIO
+# set device
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+if device.type != 'cuda':
+    raise ValueError("need to run on GPU")
+
+class EndpointHandler():
     def __init__(self, path=""):
-        self.pipeline = pipeline("text-classification",model=path)
-        self.holidays = holidays.US()
+        # load the optimized model
+        self.pipe = StableDiffusionPipeline.from_pretrained(path, torch_dtype=torch.float16)
+        self.pipe = self.pipe.to(device)
 
-    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
+
+    def __call__(self, data: Any) -> "PIL.Image":
         """
-       data args:
-            inputs (:obj: `str`)
-            date (:obj: `str`)
-      Return:
-            A :obj:`list` | `dict`: will be serialized and returned
+        Args:
+            data (:obj:):
+                includes the input data and the parameters for the inference.
+        Return:
+            A :obj:`dict`:. base64 encoded image
         """
-        # get inputs
-        inputs = data.pop("inputs",data)
-        date = data.pop("date", None)
-
-        # check if date exists and if it is holiday
-        if date is not None and date in self.holidays:
-          return [{"label": "happy", "score": 1}]
+        inputs = data.pop("inputs", data)
 
+        # run inference pipeline
+        with autocast(device.type):
+            image = self.pipe(inputs, guidance_scale=7.5)["sample"][0]
 
-        # run normal prediction
-        prediction = self.pipeline(inputs)
-        return prediction
+        # encoding image as base 64 is done by the default toolkit
+        return image