Create handler.py

handler.py

@@ -0,0 +1,24 @@
+from image_gen_aux import UpscaleWithModel
+import numpy as np
+from PIL import Image
+from io import BytesIO
+
+
+class EndpointHandler():
+    def __init__(self, path=""):
+        self.upscaler = UpscaleWithModel.from_pretrained("DAT_x3.pth").to("cuda")
+        
+    def __call__(self, data):
+        inputs = data.pop("inputs", data)
+        output = upscaler(inputs, tiling=True, tile_width=768, tile_height=768, overlap=8)
+        
+        output = outputs.reconstruction.data.squeeze().float().cpu().clamp_(0, 1).numpy()
+        output = np.moveaxis(output, source=0, destination=-1)
+        output = (output * 255.0).round().astype(np.uint8)
+
+        img = Image.fromarray(output)
+        buffered = BytesIO()
+        img.save(buffered, format="JPEG")
+        img_str = base64.b64encode(buffered.getvalue())
+
+        return img_str.decode()