|
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.float32) |
|
|
|
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": inputs.get("strength", 0.3), |
|
"guidance_scale": inputs.get("guidance_scale", 10), |
|
"height": 768, |
|
"width": 768} |
|
|
|
return self._pipe(**params).images[0] |
|
|