Update handler.py

handler.py

@@ -2,6 +2,7 @@ from image_gen_aux import UpscaleWithModel
 import numpy as np
 from PIL import Image
 from io import BytesIO
+import base64
 
 
 class EndpointHandler():
@@ -12,13 +13,8 @@ class EndpointHandler():
         img = data.pop("inputs", data)
         output = self.upscaler(img, tiling=True, tile_width=768, tile_height=768, overlap=8)
         
-        output = output.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")
+        output.save(buffered, format="JPEG")
         img_str = base64.b64encode(buffered.getvalue())
 
         return img_str.decode()