glowforge-dev/stable-diffusion-2-1-base-custom

1from typing import  Dict, List, Any2import torch3import requests4from PIL import Image5from io import BytesIO6from diffusers import StableDiffusionPipeline, StableDiffusionImg2ImgPipeline, DDIMScheduler7 8# set device9device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')10 11if device.type != 'cuda':12    raise ValueError("need to run on GPU")13 14model_id = "stabilityai/stable-diffusion-2-1-base"15 16class EndpointHandler():17    def __init__(self, path=""):18        # load the optimized model19        self.textPipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)20        self.textPipe.scheduler = DDIMScheduler.from_config(self.textPipe.scheduler.config)21        self.textPipe = self.textPipe.to(device)22 23        # create an img2img model24        self.imgPipe = StableDiffusionImg2ImgPipeline.from_pretrained(model_id, torch_dtype=torch.float16)25        self.imgPipe.scheduler = DDIMScheduler.from_config(self.imgPipe.scheduler.config)26        self.imgPipe = self.imgPipe.to(device)27 28    def __call__(self, data: Any) -> List[List[Dict[str, float]]]:29        """30        Args:31            data (:obj:):32                includes the input data and the parameters for the inference.33        Return:34            A :obj:`dict`:. base64 encoded image35        """36        prompt = data.pop("inputs", data)37        url = data.pop("url", data)38        response = requests.get(url)39        init_image = Image.open(BytesIO(response.content)).convert("RGB")40        init_image.thumbnail((512, 512))41 42        params = data.pop("parameters", data)43 44        # hyperparamters45        num_inference_steps = params.pop("num_inference_steps", 25)46        guidance_scale = params.pop("guidance_scale", 7.5)47        negative_prompt = params.pop("negative_prompt", None)48        height = params.pop("height", None)49        width = params.pop("width", None)50        manual_seed = params.pop("manual_seed", -1)51 52        out = None53 54        if data.get("url"):55            generator = torch.Generator(device='cuda')56            generator.manual_seed(manual_seed)57            # run img2img pipeline58            out = self.imgPipe(prompt, 59                        image=init_image,60                        num_inference_steps=num_inference_steps,61                        guidance_scale=guidance_scale,62                        num_images_per_prompt=1,63                        negative_prompt=negative_prompt,64                        height=height,65                        width=width66            )67        else:68            # run text pipeline69            out = self.textPipe(prompt, 70                        num_inference_steps=num_inference_steps,71                        guidance_scale=guidance_scale,72                        num_images_per_prompt=1,73                        negative_prompt=negative_prompt,74                        height=height,75                        width=width76            )77 78 79        # return first generated PIL image80        return out.images[0]81