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

def inpaint_with_mask(image: Image.Image, mask: Image.Image, prompt: str = "a scenic landscape") -> Image.Image:
    image = image.resize((512, 512))
    mask = mask.resize((512, 512)).convert("L")

    # 🧠 Load Inpainting Pipeline (SD 1.5-based, CPU compatible)
    pipe = StableDiffusionInpaintPipeline.from_pretrained(
        "runwayml/stable-diffusion-inpainting",
        torch_dtype=torch.float32
    ).to("cpu")

    # 🖌️ Inpaint
    result = pipe(prompt=prompt, image=image, mask_image=mask).images[0]

    # 🧹 Unload model to free memory
    del pipe
    gc.collect()

    return result

# 🎛️ Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("## 🎨 Inpaint with Stable Diffusion (CPU Safe – SD 1.5)")
    with gr.Row():
        image_input = gr.Image(label="Original Image", type="pil")
        mask_input = gr.Image(label="Mask (white = inpaint)", type="pil")
    prompt_input = gr.Textbox(label="Prompt", value="a scenic landscape")
    output = gr.Image(label="Inpainted Output")

    run_btn = gr.Button("Inpaint")

    run_btn.click(fn=inpaint_with_mask, inputs=[image_input, mask_input, prompt_input], outputs=output)

demo.launch(show_error=True)