File size: 608 Bytes
6862be6 7ff2823 6862be6 98fcb20 6862be6 8d69278 b2f9547 6862be6 7ff2823 6862be6 |
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_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() |