import gradio as gr
from PIL import Image
import numpy as np
from diffusers import StableDiffusionInpaintPipeline
import torch
from diffusers.utils import load_image
from diffusers.pipelines.stable_diffusion import safety_checker
def sc(self, clip_input, images) :
    return images, [False for i in images]
safety_checker.StableDiffusionSafetyChecker.forward = sc


pipe = StableDiffusionInpaintPipeline.from_pretrained(
    "runwayml/stable-diffusion-inpainting",
    revision="fp16",
    torch_dtype=torch.float16,
)
pipe.to('cpu')

def inpaint_image(image, mask, prompt, negative_prompt):
    print(image_file, mask_file, prompt, negative_prompt)
    n_image = pipe(prompt, image=image, mask_image=mask, guidance_scale=5,height=int(8*round(image.height/8)), width=int(8*round(image.width/8)), num_inference_steps=70,negative_prompt=negative_prompt).images[0]  
    return n_image

def process_files(image_file, mask_file, prompt, negative_prompt):
    image = Image.open(image_file)
    mask = Image.open(mask_file)
    return inpaint_image(image, mask, prompt, negative_prompt)

with gr.Blocks() as demo:
    gr.Markdown("## Inpainting App")

    with gr.Row():
        with gr.Column():
            image_input = gr.File(label="Input Image", type="filepath")
            mask_input = gr.File(label="Mask Image", type="filepath")
            prompt_input = gr.Textbox(label="Prompt", placeholder="Enter your prompt here...")
            negative_prompt_input = gr.Textbox(label="Negative Prompt", placeholder="Enter your negative prompt here...")

            submit_button = gr.Button("Inpaint")

        with gr.Column():
            output_image = gr.Image(type="pil", label="Inpainted Image")

    submit_button.click(
        fn=process_files,
        inputs=[image_input, mask_input, prompt_input, negative_prompt_input],
        outputs=output_image
    )

# Launch the interface
demo.launch()