| import torch | |
| from diffusers import StableDiffusionInstructPix2PixPipeline, EulerAncestralDiscreteScheduler | |
| import base64 | |
| from io import BytesIO | |
| from PIL import Image | |
| import json | |
| class EndpointHandler(): | |
| def __init__(self, path=""): | |
| model_id = "timbrooks/instruct-pix2pix" | |
| self.pipe = StableDiffusionInstructPix2PixPipeline.from_pretrained(model_id, torch_dtype=torch.float16, safety_checker=None) | |
| self.pipe.to("cuda") | |
| self.pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(self.pipe.scheduler.config) | |
| def __call__(self, data): | |
| info=data['inputs'] | |
| image=info.pop("image",data) | |
| prompt=info.pop("prompt",data) | |
| image=base64.b64decode(image) | |
| raw_images = Image.open(BytesIO(image)).convert('RGB') | |
| images = self.pipe(prompt, image=raw_images, num_inference_steps=25, image_guidance_scale=1).images | |
| return json.dumps({'image':images[0]}) | |
| if __name__=="__main__": | |
| my_handler=EndpointHandler(path='.') | |