import gradio as gr
import torch
from diffusers import StableDiffusionInpaintPipeline
from PIL import Image

model_id = "stabilityai/stable-diffusion-2-inpainting"
pipeline = StableDiffusionInpaintPipeline.from_pretrained(model_id, torch_dtype=torch.float16).to("cuda")

def inpaint_fn(image: Image.Image, mask: Image.Image, prompt: str):
    result = pipeline(prompt=prompt, image=image, mask_image=mask).images[0]
    return result