| from diffusers import AutoPipelineForImage2Image | |
| import torch | |
| from typing import Dict, Any | |
| from PIL import Image | |
| from io import BytesIO | |
| import base64 | |
| class EndpointHandler(): | |
| def __init__(self, path="."): | |
| if torch.cuda.is_available(): | |
| device = "cuda" | |
| else: | |
| device = "cpu" | |
| self._pipe = AutoPipelineForImage2Image.from_pretrained(path, torch_dtype=torch.float16).to(device) | |
| def __call__(self, data: Dict[str, Any]) -> list[Dict[str, Any]]: | |
| inputs = data.pop("inputs", data) | |
| params = {"prompt": inputs.get("prompt", ""), | |
| "image": Image.open(BytesIO(base64.b64decode(inputs['image']))), | |
| "strength": float(inputs.get("strength", 0.3)), | |
| "guidance_scale": float(inputs.get("guidance_scale", 10)), | |
| "height": 768, | |
| "width": 768} | |
| img: Image = self._pipe(**params).images[0] | |
| stream = BytesIO() | |
| img.save(stream, format="jpeg") | |
| res = {"status": 200, | |
| "image": base64.b64encode(stream.getvalue()).decode("utf8") | |
| } | |
| return res | |
| if __name__ == "__main__": | |
| h = EndpointHandler() | |
| v = h({}) | |
| print(v) |