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 generate(image_editor):
    image = image_editor['background'].convert('RGB')
    mask = Image.new("RGBA", image_editor["layers"][0].size, "WHITE") 
    mask.paste(image_editor["layers"][0], (0, 0), image_editor["layers"][0])
    mask = ImageOps.invert(mask.convert('L'))
    
    image.thumbnail((1024, 1024))
    mask.thumbnail((1024, 1024))

    return image_editor, image, mask

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')
            generate_button = gr.Button("Generate")
        with gr.Column():
            version_gallery = gr.Gallery(label="Versions")
            restore_button = gr.Button("Restore Version")

    with gr.Accordion("Advanced Settings", open=False):
        neg_prompt = gr.Textbox(label='Negative Prompt', value='ugly, deformed, nsfw')
        strength_slider = gr.Slider(0.0, 1.0, value=1.0, label="Strength")
        guidance_slider = gr.Slider(1.0, 15.0, value=7.5, label="Guidance")
        
    with gr.Row():
        out1 = gr.Image()
        out2 = gr.Image()
        out3 = gr.Image()
        
    generate_button.click(fn=generate, inputs=sketch_pad, outputs=[sketch_pad, out1, out2])

demo.launch()