QUBUHUB/stable-diffusion-v1-4-endpoints

1from typing import  Dict, List, Any2import torch3from torch import autocast4from diffusers import StableDiffusionPipeline5import base646from io import BytesIO7 8 9# set device10device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')11 12if device.type != 'cuda':13    raise ValueError("need to run on GPU")14 15class EndpointHandler():16    def __init__(self, path=""):17        # load the optimized model18        self.pipe = StableDiffusionPipeline.from_pretrained(path, torch_dtype=torch.float16)19        self.pipe = self.pipe.to(device)20 21 22    def __call__(self, data: Any) -> "PIL.Image":23        """24        Args:25            data (:obj:):26                includes the input data and the parameters for the inference.27        Return:28            A :obj:`dict`:. base64 encoded image29        """30        inputs = data.pop("inputs", data)31        32        # run inference pipeline33        with autocast(device.type):34            image = self.pipe(inputs, guidance_scale=7.5)["sample"][0]  35            36        # encoding image as base 64 is done by the default toolkit37        return image38