Update handler.py

handler.py

@@ -23,9 +23,9 @@ class EndpointHandler():
     def __init__(self, path=""):
         self.path = path
         # load the optimized model
-        #model_id = "stabilityai/stable-diffusion-x4-upscaler"
-        #self.pipe = StableDiffusionPipeline.from_pretrained(path, torch_dtype=torch.float16)
-        #self.pipe = self.pipe.to(device)
+        model_id = "stabilityai/stable-diffusion-x4-upscaler"
+        self.pipe = StableDiffusionPipeline.from_pretrained(path, torch_dtype=torch.float16)
+        self.pipe = self.pipe.to(device)
 
 
     def __call__(self, data) -> List[Dict[str, Any]]:
@@ -39,5 +39,13 @@ class EndpointHandler():
         inputs = data.get("inputs")
         logger.info('inputs received %s', inputs)
 
+        decoded_image = Image.open(BytesIO(base64.b64decode(inputs['image'])))
+        with autocast(device.type):
+            upscaled_image = self.pipe(prompt="", image = decoded_image).images[0]
+
+        buffered = BytesIO()
+        upscaled_image.save(buffered, format="JPEG")
+        img_str = base64.b64encode(buffered.getvalue())
+
         # postprocess the prediction
-        return {"image": "nothing"}
+        return {"image": img_str}