import gradio as gr
import numpy as np
import torch
from diffusers import StableDiffusionInpaintPipeline
from PIL import Image
from segment_anything import SamPredictor, sam_model_registry

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
sam_checkpoint = "sam_vit_h_4b8939.pth"
model_type = "vit_h"
sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
sam.to(device)
predictor = SamPredictor(sam)

pipe = StableDiffusionInpaintPipeline.from_pretrained(
    "stabilityai/stable-diffusion-2-inpainting",
    # torch_dtype = torch.float16
)
pipe = pipe.to(device)

selected_pixels = []

with gr.Blocks() as demo:
    with gr.Row():
        input_img = gr.Image(label="Input")
        mask_img = gr.Image(label="Mask")
        output_img = gr.Image(label="Output")
    
    with gr.Blocks():
        prompt_text = gr.Textbox(lines=1, label="Prompt")

    with gr.Row():
        submit =gr.Button("submit")

    def inpaint(image, mask, prompt):
        image = Image.fromarray(image)
        mask = Image.fromarray(mask)

        image = image.resize((512,512))
        mask = mask.resize((512,512))

        output = pipe(
            prompt = prompt,
            image= image,
            mask_image = mask
        ).images[0]

        return output

    def generate_mask(image, evt:gr.SelectData):
        selected_pixels.append(evt.index)
        predictor.set_image(image)
        input_points = np.array(selected_pixels)
        input_labels = np.ones(input_points.shape[0])
        mask, _, _ = predictor.predict(
            point_coords=input_points,
            point_labels=input_labels,
            multimask_output=False
        )
        #(n,size,size)
        # mask = np.logical_not(mask) # for selection of foregraound and background will change.
        mask = Image.fromarray(mask[0, :, :])      
        return mask
    
    input_img.select(generate_mask, [input_img],[mask_img])
    submit.click(inpaint, inputs=[input_img,mask_img,prompt_text], outputs=[output_img])

if __name__ == "__main__":
    demo.launch()