File size: 2,212 Bytes
7dbdbc9 3c71e14 b6b2b75 3c71e14 7dbdbc9 63f6eb9 7dbdbc9 3c71e14 7dbdbc9 3c71e14 e51059c 3c71e14 e51059c 3c71e14 e51059c 082ff75 3c71e14 e51059c 3c71e14 b54e214 |
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 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 |
import base64
from io import BytesIO
from typing import Any, Dict
import torch
from diffusers import FluxFillPipeline
from PIL import Image
def decode_image(b64_string):
image_data = base64.b64decode(b64_string)
return Image.open(BytesIO(image_data)).convert("RGB")
def encode_image(image):
buffer = BytesIO()
image.save(buffer, format="PNG")
return base64.b64encode(buffer.getvalue()).decode("utf-8")
class EndpointHandler:
def __init__(self, path="shangguanyanyan/flux1-fill-dev-custom"):
self.pipe = FluxFillPipeline.from_pretrained(
path, torch_dtype=torch.bfloat16
).to("cuda" if torch.cuda.is_available() else "cpu")
self.parameters = {
"height": 1632,
"width": 1232,
"guidance_scale": 30,
"num_inference_steps": 50,
"max_sequence_length": 512,
"generator": torch.Generator("cpu").manual_seed(0),
}
def __call__(self, data: Any) -> Dict[str, Any]:
"""
data: {
"inputs": {
"image": base64_image,
"mask": base64_mask,
"prompt": prompt
},
"parameters": {
"height": 1632,
"width": 1232,
"guidance_scale": 30,
"num_inference_steps": 50,
"max_sequence_length": 512,
}
}
"""
inputs = data.pop("inputs", data)
parameters = data.pop("parameters", {})
parameters.update(self.parameters)
base64_image = inputs.pop("image", "")
base64_mask = inputs.pop("mask", "")
prompt = inputs.pop("prompt", "")
if not base64_image or not base64_mask or not prompt:
return {
"error": "Please provide image, mask and prompt",
"status": "failed",
}
image = decode_image(base64_image)
mask = decode_image(base64_mask)
image = self.pipe(
prompt=prompt,
image=image,
mask_image=mask,
**parameters,
).images[0]
return {"image": encode_image(image), "status": "success"}
|