import gradio as gr
from diffusers import AutoPipelineForInpainting, AutoencoderKL
import torch
from PIL import Image, ImageOps

vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
pipeline = AutoPipelineForInpainting.from_pretrained("diffusers/stable-diffusion-xl-1.0-inpainting-0.1", vae=vae, torch_dtype=torch.float16, variant="fp16", use_safetensors=True).to("cuda")

def get_select_index(evt: gr.SelectData):
    return evt.index

def squarify_image(img):
    if(img.height > img.width): bg_size = img.height
    else:  bg_size = img.width
    bg = Image.new(mode="RGB", size=(bg_size,bg_size), color="white")
    bg.paste(img, ( int((bg.width - bg.width)/2), 0) )

    return bg
    
def divisible_by_8(image):
    width, height = image.size
    
    # Calculate the new width and height that are divisible by 8
    new_width = (width // 8) * 8
    new_height = (height // 8) * 8
    
    # Resize the image
    resized_image = image.resize((new_width, new_height))
    
    return resized_image


def restore_version(index, versions):
    final_dict = {'background': versions[index], 'layers': None, 'composite': versions[index]}
    return final_dict


def generate(image_editor, prompt, neg_prompt, versions):
    print('versions')
    print(versions)
    
    image = image_editor['background'].convert('RGB')
    image.thumbnail((1024, 1024))
    image = divisible_by_8(image)
    original_image_size = image.size
    
    layer = image_editor["layers"][0].resize(image.size)
    
    image = squarify_image(image)
    
    mask = Image.new("RGBA", image.size, "WHITE") 
    mask.paste(layer, (0, 0), layer)
    mask = ImageOps.invert(mask.convert('L'))

    final_image = pipeline(prompt=prompt, 
                           image=image, 
                           mask_image=mask).images[0]
    
    final_image = final_image.crop((0, 0, original_image_size[0], original_image_size[1]))

    final_dict = {'background': final_image, 'layers': None, 'composite': final_image}

    if(versions==None): 
        final_gallery = [image_editor['background'] ,final_image]
    else: 
        final_gallery = versions
        final_gallery.append(final_image)
    
    return final_dict, image, mask, final_image, final_gallery

with gr.Blocks() as demo:
    gr.Markdown("""
    # Inpainting Sketch Pad
    by [Tony Assi](https://www.tonyassi.com/)
    """)
    
    with gr.Row():
        with gr.Column():
            sketch_pad = gr.ImageMask(type='pil', label='Inpaint')
            prompt = gr.Textbox(label="Prompt")
            generate_button = gr.Button("Generate")
            with gr.Accordion("Advanced Settings", open=False):
                neg_prompt = gr.Textbox(label='Negative Prompt', value='ugly, deformed')
        with gr.Column():
            version_gallery = gr.Gallery(label="Versions", type="pil", object_fit='contain')
            restore_button = gr.Button("Restore Version")
            selected = gr.Number(show_label=False, visible=True)
            
        
    with gr.Row():
        out1 = gr.Image(format='png')
        out2 = gr.Image(format='png')
        out3 = gr.Image(format='png')

    version_gallery.select(get_select_index, None, selected)
    generate_button.click(fn=generate, inputs=[sketch_pad,prompt, neg_prompt, version_gallery], outputs=[sketch_pad, out1, out2, out3, version_gallery])
    restore_button.click(fn=restore_version, inputs=[selected, version_gallery], outputs=sketch_pad)

demo.launch()