from image_gen_aux import UpscaleWithModel
import numpy as np
from PIL import Image
from io import BytesIO
import base64


class EndpointHandler():
    def __init__(self, path=""):
        self.upscaler = UpscaleWithModel.from_pretrained("OzzyGT/DAT_X2").to("cuda")

    def __call__(self, data):
        img = data.pop("inputs", data)
        output = self.upscaler(img, tiling=True, tile_width=768, tile_height=768, overlap=8)

        buffered = BytesIO()
        output.save(buffered, format="JPEG")
        img_str = base64.b64encode(buffered.getvalue())

        return img_str.decode()