File size: 1,342 Bytes
e6b188b a712455 e6b188b bd90d6e b13e405 96b7249 bd90d6e e6b188b b13e405 e6b188b a712455 e6b188b a712455 14df57b 2a60e8f dc4454d ad067ed a712455 034b0b6 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 |
from typing import Dict, List, Any
from transformers import AutoImageProcessor, Swin2SRForImageSuperResolution
import torch
import base64
import logging
import numpy as np
from PIL import Image
from io import BytesIO
logger = logging.getLogger()
logger.setLevel(logging.DEBUG)
# check for GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
class EndpointHandler:
def __init__(self, path=""):
# load the model
self.processor = AutoImageProcessor.from_pretrained("caidas/swin2SR-classical-sr-x2-64")
self.model = Swin2SRForImageSuperResolution.from_pretrained("caidas/swin2SR-classical-sr-x2-64")
# move model to device
self.model.to(device)
def __call__(self, data: Any):
image = data["inputs"]
inputs = self.processor(image, return_tensors="pt").to(device)
with torch.no_grad():
outputs = self.model(**inputs)
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()
|