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