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

# Ensure we use GPU if available
device = "cuda" if torch.cuda.is_available() else "cpu"

def load_model():
    """Load Stable Diffusion 2 Inpainting model."""
    pipe = StableDiffusionInpaintPipeline.from_pretrained(
        "stabilityai/stable-diffusion-2-inpainting",
        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
    ).to(device)
    return pipe

pipe = load_model()

def inpaint(image_editor_output, prompt):
    """Perform inpainting using the uploaded image and mask."""
    print(f"image_editor_output:{image_editor_output}")
    
    if not isinstance(image_editor_output, dict) or "background" not in image_editor_output or "layers" not in image_editor_output:
        raise ValueError("Invalid input. Please upload an image and paint a mask.")
    
    # Extract the base image (background or composite)
    image = Image.fromarray(image_editor_output["composite"]).convert("RGB")  # Ensure image is in RGB format
    
    # Extract the first layer (mask) and convert to grayscale
    mask_array = np.array(image_editor_output["layers"][0], dtype=np.uint8)
    mask = Image.fromarray(mask_array[:, :, 3]).convert("L")  # Use the alpha channel
    
    # Resize mask to match image size if needed
    if mask.size != image.size:
        mask = mask.resize(image.size, Image.LANCZOS)
    
    # Perform inpainting
    result = pipe(prompt=prompt, image=image, mask_image=mask).images[0]
    return result

# Gradio UI with ImageEditor for inpainting
demo = gr.Interface(
    fn=inpaint,
    inputs=[
        gr.ImageEditor(label="Upload & Paint Image (White areas will be inpainted)"),
        gr.Textbox(label="Prompt for Inpainting"),
    ],
    outputs=gr.Image(type="pil", label="Generated Image"),
    title="Stable Diffusion 2 Inpainting",
    description="Upload an image, paint directly on it to define mask areas, and provide a text prompt.",
)

demo.launch()