Create handler.py

handler.py

@@ -0,0 +1,29 @@
+from diffusers import DiffusionPipeline
+
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+if device.type != 'cuda':
+    raise ValueError("need to run on GPU")
+# set mixed precision dtype
+dtype = torch.bfloat16 if torch.cuda.get_device_capability()[0] == 8 else torch.float16
+
+class EndpointHandler():
+    def __init__(self, path=""):
+        self.pipeline = DiffusionPipeline.from_pretrained("CompVis/ldm-super-resolution-4x-openimages", torch_dtype=dtype).to(device)
+        
+
+    def __call__(self, data: Any) -> List[List[Dict[str, float]]]:
+        image = data.pop("image", None)
+        low_res_img = image.resize((128, 128))
+        
+        with torch.no_grad():
+            upscaled_image = self.pipeline(low_res_img, num_inference_steps=100, eta=1).images[0]
+        
+        return upscaled_image
+
+        
+    # helper to decode input image
+    def decode_base64_image(self, image_string):
+        base64_image = base64.b64decode(image_string)
+        buffer = BytesIO(base64_image)
+        image = Image.open(buffer)
+        return image