File size: 608 Bytes
52bddc6 0b005f7 52bddc6 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 |
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_X4").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() |