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

# Load the Stable Diffusion inpainting model (ensure GPU usage if available)
pipe = StableDiffusionInpaintPipeline.from_pretrained(
    "stabilityai/stable-diffusion-2-1",  # You can use another inpainting model here
    torch_dtype=torch.float16
).to("cuda")

# Function to remove makeup based on prompt and image
def remove_makeup(image, prompt):
    # Debugging: Check the type of the input image
    print(f"Input image type: {type(image)}")
    
    # Ensure the input image is a PIL Image
    if not isinstance(image, Image.Image):
        image = Image.fromarray(image)
    
    # Debugging: Print the size before resizing
    print(f"Original image size: {image.size}")

    # Resize the image to 512x512 (standard input size for Stable Diffusion)
    image = image.resize((512, 512))

    # Debugging: Print the size after resizing
    print(f"Resized image size: {image.size}")

    # Convert the image to a numpy array and normalize to range [-1, 1]
    image_array = np.array(image) / 255.0
    image_array = (image_array - 0.5) / 0.5  # Normalize to [-1, 1]
    
    # Debugging: Check the shape and type of the image array
    print(f"Image array shape: {image_array.shape}, type: {type(image_array)}")

    # Convert to tensor and move to GPU
    image_tensor = torch.tensor(image_array).permute(2, 0, 1).unsqueeze(0).float().to("cuda")  # Convert to tensor

    # Debugging: Check the shape and type of the tensor before sending it to the model
    print(f"Image tensor shape: {image_tensor.shape}, type: {type(image_tensor)}")

    # Use the model to process the image with the given prompt
    result = pipe(prompt=prompt, image=image_tensor).images[0]

    # Return the result
    return result

# Define the Gradio interface
iface = gr.Interface(
    fn=remove_makeup,
    inputs=[ 
        gr.Image(type="pil", label="Upload Image with Makeup"),  # Input image (PIL)
        gr.Textbox(label="Prompt", placeholder="Describe what you want (e.g., 'a natural face with no makeup and no filter')")  # Input prompt
    ],
    outputs=gr.Image(type="pil", label="Makeup Removed Image"),
    title="Makeup Removal Using Stable Diffusion",
    description="Upload a face image with makeup and provide a prompt (e.g., 'a natural face with no makeup and no filter'). The model will remove makeup and filters, returning a natural-looking face.",
)

# Launch the Gradio interface
iface.launch()