from typing import Dict, List, Any
import torch
from torch import autocast
from diffusers import StableDiffusionAdapterPipeline
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 = StableDiffusionAdapterPipeline.from_pretrained(path, torch_dtype=torch.float16)
    self.pipe = self.pipe.to(device)
  
  
  def __call__(self, data:Any) -> List[List[Dict[str, float]]]:
    inputs = data.pop("inputs", data)
    
    with autocast(device.type):
      image = self.pipe(inputs, guidance_scale=7.5)["samples"][0]
      
    buffered = BytesIO()
    image.save(buffered, format="JPEG")
    img_str = base64.base64encode(buffered.getvalue())
    
    return { "image" : img_str.decode()}