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