| 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() |