import gradio as gr
from PIL import Image

import torch
from diffusers import AutoPipelineForInpainting
from diffusers.utils import load_image

def draw_on_image(image, prompt):
    print(image, prompt)

    if not prompt:
        return

    init_image = load_image(
        image["image"]
    )
    mask_image = load_image(
        image["mask"]
    )

    res_image = pipeline(prompt=prompt, image=init_image, mask_image=mask_image).images[0]

    return res_image

inputs = [
            gr.Image(tool="sketch", label="Image", type="pil"),
            gr.Text(max_lines=1)
    ]

if torch.cuda.is_available():
    torch_dtype = torch.float32
    device = "cuda"
else:
    torch_dtype = torch.float16
    device = "cpu"

pipeline = AutoPipelineForInpainting.from_pretrained(
    "kandinsky-community/kandinsky-2-2-decoder-inpaint", torch_dtype=torch_dtype
)

pipeline.to(device)

app = gr.Interface(draw_on_image, inputs=inputs, outputs="image")
app.queue()
app.launch(share=True)