BuildingVirtuallyAI/RenderingModel

1from typing import  Dict, Any2import torch3import base644from io import BytesIO5from model import Model6from PIL import Image7# set device8device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')9 10if device.type != 'cuda':11    raise ValueError("need to run on GPU")12 13class EndpointHandler():14    def __init__(self, path=""):15        # load the optimized model16        self.model = Model()17 18 19    def __call__(self, data: Any) -> Any:20        """21        Args:22            data (:obj:):23                includes the input data and the parameters for the inference.24        Return:25            A :obj:`dict`:. base64 encoded image26        """27        inputs = data.pop("inputs", data)28 29        30        image = Image.open(BytesIO(base64.b64decode(inputs['image'])))31 32        # run inference pipeline33        _, res = self.model.process_lineart(image) 34        35            36        # encoding image as base 64 is done by the default toolkit37        return res38