from typing import  Dict, List, Any
import torch
from diffusers import StableDiffusionXLImg2ImgPipeline, DiffusionPipeline, AutoencoderKL, DPMSolverMultistepScheduler, DDIMScheduler, StableDiffusionInpaintPipeline, AutoPipelineForInpainting, AutoPipelineForImage2Image, StableDiffusionControlNetInpaintPipeline, ControlNetModel
from PIL import Image
import base64
from io import BytesIO
from diffusers.image_processor import VaeImageProcessor
import numpy as np


# set device
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.smooth_pipe = StableDiffusionXLImg2ImgPipeline.from_pretrained(
          "stabilityai/stable-diffusion-xl-refiner-1.0", torch_dtype=torch.float16
        )
        self.smooth_pipe.to("cuda")

        self.vae = AutoencoderKL.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0",
                                                 subfolder="vae",  use_safetensors=True,
                                                 ).to("cuda")

        self.smooth_pipe.enable_model_cpu_offload()
        self.smooth_pipe.enable_xformers_memory_efficient_attention()


        self.controlnet = ControlNetModel.from_pretrained(
            "lllyasviel/control_v11p_sd15_inpaint", torch_dtype=torch.float16
        )

        self.pipe = StableDiffusionControlNetInpaintPipeline.from_pretrained(
            "runwayml/stable-diffusion-v1-5", controlnet=self.controlnet, torch_dtype=torch.float16
        )

        self.pipe.scheduler = DDIMScheduler.from_config(self.pipe.scheduler.config)
        self.pipe.enable_model_cpu_offload()
        self.pipe.enable_xformers_memory_efficient_attention()
        


    def __call__(self, data: Any) -> List[List[Dict[str, float]]]:
        """
        :param data: A dictionary contains `inputs` and optional `image` field.
        :return: A dictionary with `image` field contains image in base64.
        """

        method = data.pop("method", "rasterize")

        if(method == "rasterize"):
        
            encoded_image = data.pop("image", None)
            
            prompt = data.pop("prompt", "")
            num_inference_steps = data.pop("num_inference_steps", 50)
    
            if encoded_image is not None:
                image = self.decode_base64_image(encoded_image).convert('RGB')
    
                image_processor = VaeImageProcessor();
                latents = image_processor.preprocess(image)
                latents = latents.to(device="cuda")
                
                with torch.no_grad():
                    latents_dist = self.vae.encode(latents).latent_dist.sample() * self.vae.config.scaling_factor
                
                self.smooth_pipe.enable_xformers_memory_efficient_attention()
                out = self.smooth_pipe(prompt, image=latents_dist, num_inference_steps=num_inference_steps).images
    
                return out
        else:
            encoded_image = data.pop("image", None)
            encoded_mask_image = data.pop("mask_image", None)
            
            prompt = data.pop("prompt", "")
    
            negative_prompt = data.pop("negative_prompt", "")
    
            method = data.pop("method", "slow")
            strength = data.pop("strength", 0.2)
            guidance_scale = data.pop("guidance_scale", 8.0)
            num_inference_steps = data.pop("num_inference_steps", 20)

            # process image
            if encoded_image is not None and encoded_mask_image is not None:
                image = self.decode_base64_image(encoded_image).convert("RGB")
                mask_image = self.decode_base64_image(encoded_mask_image).convert("RGB")
            else:
                image = None
                mask_image = None

            control_image = self.make_inpaint_condition(image, mask_image)

            # generate image
            image = self.pipe(
                prompt=prompt,
                negative_prompt=negative_prompt,
                num_inference_steps=num_inference_steps,
                eta=1.0,
                image=image,
                mask_image=mask_image,
                control_image=control_image,
                guidance_scale=guidance_scale,
                strength=strength
            ).images[0]
    
            return image
    
    # helper to decode input image
    def decode_base64_image(self, image_string):
        base64_image = base64.b64decode(image_string)
        buffer = BytesIO(base64_image)
        image = Image.open(buffer)
        return image
        
    def make_inpaint_condition(self, image, image_mask):
        image = np.array(image.convert("RGB")).astype(np.float32) / 255.0
        image_mask = np.array(image_mask.convert("L")).astype(np.float32) / 255.0
    
        assert image.shape[0:1] == image_mask.shape[0:1], "image and image_mask must have the same image size"
        image[image_mask > 0.5] = -1.0  # set as masked pixel
        image = np.expand_dims(image, 0).transpose(0, 3, 1, 2)
        image = torch.from_numpy(image)
        return image