File size: 1,198 Bytes
9701c1e 9b5b006 9701c1e c968655 335f64a 9701c1e 335f64a cebbecd 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 31 32 33 34 |
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("text",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
img=images[0]
img.save("./1.png")
with open('./1.png','rb') as img_file:
encoded_string = base64.b64encode(img_file.read()).decode('utf-8')
return {'image':encoded_string}
if __name__=="__main__":
my_handler=EndpointHandler(path='.')
|