from typing import Dict, List, Any from transformers import Pipeline from PIL import Image from io import BytesIO import base64 import json from visual_chatgpt import ImageEditing, Text2Box, Segmenting, Inpainting import os class EndpointHandler(): def __init__(self, path=""): # 换个顺序就行了?????离谱 os.system("ln -s /repository /app/repository") self.sam = Segmenting('cuda') self.inpaint = Inpainting('cuda') self.grounding = Text2Box('cuda') self.model = ImageEditing(self.grounding,self.sam,self.inpaint) def __call__(self, data): # data=json.loads(data) # inputs=data.pop("inputs",data) # inputs=base64.b64decode(inputs) # raw_images = Image.open(BytesIO(inputs)) info=data['inputs'] image=info.pop('image',data) image=base64.b64decode(image) raw_image=Image.open(BytesIO(image)).convert('RGB') texts=info.pop('texts',data) target=texts[0] replacement=texts[1] if replacement=="": img=self.model.inference_remove(raw_image,target) img_data=img.tobytes() img_base64=base64.b64encode(img_data).decode('utf-8') return {'image':img_base64} else: img=self.model.inference_replace_sam(raw_image,target,replacement) img_data=img.tobytes() img_base64=base64.b64encode(img_data).decode('utf-8') return {'image':img_base64} if __name__=="__main__": my_handler=EndpointHandler(path='.') # test_payload={"inputs": "/home/ubuntu/guoling/1.png"} with open("/home/ubuntu/guoling/1.png",'rb') as img: image_bytes=img.read() image_base64=base64.b64encode(image_bytes).decode('utf-8') target="the pig" replacement="" data={ 'inputs':{ "image":image_base64, "target":target, "replacement":replacement } } result=my_handler(data) result.save("new1.png")