|
from typing import Dict, List, Any |
|
from PIL import Image |
|
from io import BytesIO |
|
import torch |
|
import base64 |
|
from diffusers import StableDiffusionInstructPix2PixPipeline, EulerAncestralDiscreteScheduler |
|
import diffusers |
|
import transformers |
|
import logging |
|
import subprocess |
|
import sys |
|
|
|
logger = logging.getLogger() |
|
logger.setLevel(logging.DEBUG) |
|
|
|
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') |
|
subprocess.run("nvidia-smi") |
|
|
|
logger.info(f"torch version: {torch.__version__}") |
|
logger.info(f"diffusers version: {diffusers.__version__}") |
|
logger.info(f"transformers version: {transformers.__version__}") |
|
|
|
logger.info(f"device: {device}") |
|
|
|
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.scheduler = EulerAncestralDiscreteScheduler.from_config(self.pipe.scheduler.config) |
|
self.pipe.to(device) |
|
|
|
logger.info(f"PIPE LOADED") |
|
|
|
def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]: |
|
""" |
|
data dict: |
|
inputs: base64 encoded image, |
|
parameters: dict: |
|
prompt: str |
|
returns: |
|
base64 encoded image |
|
""" |
|
|
|
|
|
image_data = data.pop("inputs", data) |
|
logger.info(f"Raw img size: {sys.getsizeof(image_data)}") |
|
|
|
image = Image.open(BytesIO(base64.b64decode(image_data))) |
|
logger.info(f"PIL Image img size: {sys.getsizeof(image)}") |
|
|
|
parameters = data.pop("parameters", data) |
|
prompt = parameters['prompt'] |
|
|
|
images = self.pipe(prompt, image=image, num_inference_steps=10, image_guidance_scale=1).images |
|
|
|
return images[0] |
