| from typing import Dict, List, Any | |
| import torch | |
| from torch import autocast | |
| from diffusers import StableDiffusionPipeline | |
| import base64 | |
| from io import BytesIO | |
| device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') | |
| if device.type != 'cuda': | |
| raise ValueError("Must run SDXL on a GPU instance.") | |
| class EndpointHandler(): | |
| def __init__(self,path=""): | |
| self.pipe = StableDiffusionPipeline.from_pretrained(path,torch_dtype=torch.float16) | |
| self.pipe = self.pipe.to(device) | |
| def __call__(self): | |
| """ | |
| """ | |
| inputs = data.pop("inputs",data) | |
| with autocast(device.type): | |
| image = self.pipe(inputs,guidance_scale=9)["sample"][0] | |
| buffer = BytesIO() | |
| image.save(buffer, format="JPEG") | |
| img_str = base64.b64decode(buffer.getvalue()) | |
| return {"image": img_str.decode} |