Update handler.py

handler.py

@@ -10,13 +10,12 @@ from realesrgan import RealESRGANer
 
 class EndpointHandler:
     def __init__(self, path=""):
-        model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
+        model = RRDBNet(num_in_ch=3, num_out_ch=3)
         self.upsampler = RealESRGANer(
             scale=4,
             model_path=str(Path(path) / "RealESRGAN_x4plus.pth"),
             model=model,
             tile=0,
-            tile_pad=10,
             pre_pad=0,
             half=True,
         )
@@ -30,18 +29,21 @@ class EndpointHandler:
             A :obj:`dict`:. base64 encoded image
         """
         image = data.pop("inputs", data)
-        # image = Image.open(BytesIO(image)).convert("RGB")
+
+        # This lets us pass local images as well while developing
+        if isinstance(image, bytes):
+            image = Image.open(BytesIO(image)) # .convert("RGB")
+
         image = np.array(image)
         image = image[:, :, ::-1]  # RGB -> BGR
-
         image, _ = self.upsampler.enhance(image, outscale=4)
         image = image[:, :, ::-1]  # BGR -> RGB
         image = Image.fromarray(image)
 
-        # encode image as base 64
+        # Turn output image into bytestr
         buffered = BytesIO()
         image.save(buffered, format="PNG")
-        img_str = b64encode(buffered.getvalue())
+        img_bytes = b64encode(buffered.getvalue())
+        img_str = img_bytes.decode()
 
-        # postprocess the prediction
-        return {"image": img_str.decode()}
+        return {"image": img_str}