File size: 1,031 Bytes
9701c1e 9b5b006 9701c1e c968655 9c246cc 9701c1e 9b5b006 9701c1e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 |
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='.')
|