File size: 976 Bytes
5fb4029 49fc5ea ffa81ea 49fc5ea 5fb4029 49fc5ea 5fb4029 49fc5ea 5fb4029 ffa81ea 5fb4029 f4f42c3 3de8d46 1c337b5 5fb4029 3de8d46 15c460b 5fb4029 d46f57b 5fb4029 |
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 |
from typing import Dict, List, Any
import torch
from torch import autocast
from diffusers import StableDiffusionXLPipeline
import base64
from io import BytesIO
device = torch.device("cuda" if torch.cuda.is_available() else 'cpu')
if device.type != "cuda":
raise ValueError('need to run on gpu')
class EndpointHandler():
def __init__(self, path="") :
self.pipe = StableDiffusionXLPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16, variant="fp16", use_safetensors=True)
self.pipe = self.pipe.to(device)
def __call__(self, data:Any) -> List[List[Dict[str, float]]]:
print(data)
inputs = data.pop("inputs", data)
print(device)
with autocast(device.type):
image = self.pipe(inputs, guidance_scale=7.5).images[0]
buffered = BytesIO()
image.save(buffered, format="JPEG")
img_str = base64.b64encode(buffered.getvalue())
return { "image" : img_str.decode()}
|