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):
        img = data.pop("inputs", data)
        img = Image.open(BytesIO(img))
        output = upscaler(img, 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()