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}