|
from typing import Dict, List, Any |
|
from PIL import Image |
|
from io import BytesIO |
|
import torch |
|
import base64 |
|
from diffusers import StableDiffusionInstructPix2PixPipeline, EulerAncestralDiscreteScheduler |
|
|
|
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') |
|
|
|
class EndpointHandler(): |
|
def __init__(self, path=""): |
|
model_id = "timbrooks/instruct-pix2pix" |
|
self.pipe = StableDiffusionInstructPix2PixPipeline.from_pretrained(model_id, torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32, safety_checker=None) |
|
self.pipe.to(device) |
|
self.pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(self.pipe.scheduler.config) |
|
|
|
def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]: |
|
""" |
|
data args: |
|
inputs (:obj:`string`) |
|
parameters (:obj:) |
|
Return: |
|
A :obj:`string`:. image string |
|
""" |
|
|
|
|
|
image_data = data.pop('inputs', data) |
|
|
|
image = Image.open(BytesIO(base64.b64decode(image_data))) |
|
|
|
parameters = data.pop('parameters', []) |
|
prompt = parameters.pop('prompt', None) |
|
negative_prompt = parameters.pop('negative_prompt', None) |
|
num_inference_steps = parameters.pop('num_inference_steps', 10) |
|
image_guidance_scale = parameters.pop('image_guidance_scale', 1.5) |
|
guidance_scale = parameters.pop('guidance_scale', 7.5) |
|
|
|
|
|
images = self.pipe( |
|
prompt, |
|
image = image, |
|
negative_prompt = negative_prompt, |
|
num_inference_steps = num_inference_steps, |
|
image_guidance_scale = image_guidance_scale, |
|
guidance_scale = guidance_scale |
|
).images |
|
|
|
return images[0] |
