File size: 1,466 Bytes
0dadd79 0e2a33a 71ffd5c 0dadd79 c058aa4 510acad 0e2a33a 0dadd79 71ffd5c 0dadd79 71ffd5c 0dadd79 6845182 0dadd79 b074d05 0dadd79 5373874 |
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 |
from typing import Dict, List, Any
from diffusers import StableDiffusionImg2ImgPipeline
from diffusers.utils import load_image
import base64
from io import BytesIO
from pathlib import Path
import os
from diffusers.utils import load_image
class EndpointHandler():
def __init__(self, path=""):
repo_id = "runwayml/stable-diffusion-v1-5"
self.pipeline = StableDiffusionImg2ImgPipeline.from_pretrained(repo_id).to("cuda")
weight_name = "pixel-portrait-v1.safetensors"
self.pipeline.load_lora_weights("simulationcartridge/ppl", weight_name=weight_name)
def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
"""
data args:
inputs (:obj: `str` | `PIL.Image` | `np.array`)
kwargs
Return:
A :obj:`list` | `dict`: will be serialized and returned
"""
input_image_url = data.pop("input_image", data)
prompt = data.pop("prompt", None)
#load image
input_image = load_image(input_image_url)
# run normal prediction
output = self.pipeline(prompt=prompt, image=input_image, guidance_scale=16)
image = output.images[0]
# encode image as base 64
buffered = BytesIO()
image.save(buffered, format="JPEG")
img_str = base64.b64encode(buffered.getvalue()).decode('utf-8')
# postprocess the prediction
return {"image": img_str} |