handler.py

# install thing, just like in segment anything


from typing import Dict, List, Any
from PIL import Image
from io import BytesIO
from transformers import AutoModelForSemanticSegmentation, AutoFeatureExtractor
import base64
import torch
from torch import nn


import subprocess
result = subprocess.run(["pip", "install", "git+https://github.com/sberbank-ai/Real-ESRGAN.git"], check=True)
print(f"git+https://github.com/sberbank-ai/Real-ESRGAN.git = {result}")
from RealESRGAN import RealESRGAN


class EndpointHandler():
    def __init__(self, path="."):
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.model = RealESRGAN(self.device, scale=2)
        self.model.load_weights('/repository/RealESRGAN_x2.pth', download=True) 

        
    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
        """
       data args:
            images (:obj:`PIL.Image`)
            candiates (:obj:`list`)
      Return:
            A :obj:`list`:. The list contains items that are dicts should be liked {"label": "XXX", "score": 0.82}
        """
        inputs = data.pop("inputs", data)

        # decode base64 image to PIL
        image = Image.open(BytesIO(base64.b64decode(inputs['image'])))

        # forward pass
        output_image = self.model.predict(image)

        # base64 encode output
        buffered = BytesIO()
        output_image = output_image.convert('RGB')
        output_image.save(buffered, format="png")
        img_str = base64.b64encode(buffered.getvalue())

        # postprocess the prediction
        return {"image": img_str.decode()}