| from typing import Dict, List, Any | |
| import sys | |
| import base64 | |
| import logging | |
| import copy | |
| import numpy as np | |
| import torch | |
| from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler | |
| class EndpointHandler(): | |
| def __init__(self, path="", model_base="runwayml/stable-diffusion-v1-5"): | |
| self.pipe = StableDiffusionPipeline.from_pretrained(model_base, torch_dtype=torch.float16) | |
| self.pipe.scheduler = DPMSolverMultistepScheduler.from_config(self.pipe.scheduler.config) | |
| self.original_unet = copy.deepcopy(self.pipe.unet) | |
| def _generate_images( | |
| self, | |
| model_path, | |
| prompt, | |
| num_inference_steps=25, | |
| guidance_scale=7.5, | |
| num_images_per_prompt=1): | |
| if model_path == "base": | |
| self.pipe.unet = copy.deepcopy(self.original_unet) | |
| else: | |
| self.pipe.unet.load_attn_procs(model_path) | |
| self.pipe.to("cuda") | |
| pil_images = self.pipe( | |
| prompt=prompt, | |
| num_inference_steps=num_inference_steps, | |
| guidance_scale=guidance_scale, | |
| num_images_per_prompt=num_images_per_prompt).images | |
| np_images = [] | |
| for i in range(len(pil_images)): | |
| np_images.append(np.asarray(pil_images[i])) | |
| return np.stack(np_images, axis=0) | |
| def __call__(self, data: Dict[str, Any]) -> str: | |
| prompt = data.pop("inputs", "test image") | |
| model_path = data.pop("model_path", "base") | |
| num_inference_steps = data.pop("num_inference_steps", 25) | |
| guidance_scale = data.pop("guidance_scale", 7.5) | |
| num_images_per_prompt = data.pop("num_images_per_prompt", 1) | |
| images = self._generate_images( | |
| model_path, prompt, | |
| num_inference_steps, guidance_scale, num_images_per_prompt | |
| ) | |
| return base64.b64encode(images.tobytes()).decode() |