import streamlit as st
from PIL import Image
import torch
from diffusers import StableDiffusionInpaintPipeline
import numpy as np

# Load the StableDiffusionInpaintPipeline
pipe = StableDiffusionInpaintPipeline.from_pretrained("stabilityai/stable-diffusion-2-inpainting")
pipe.to("cuda" if torch.cuda.is_available() else "cpu")  # Move model to GPU if available

# Function to process the image with the provided prompt
def process_image(image, prompt):
    # Ensure the image is a PIL Image
    if isinstance(image, np.ndarray):
        image = Image.fromarray(image)  # Convert numpy array to PIL Image
    elif isinstance(image, torch.Tensor):
        image = Image.fromarray(image.numpy())  # Convert torch tensor to PIL Image
    
    # Resize image to the required size (512x512)
    image = image.resize((512, 512))

    # Process the image through the pipeline
    edited_image = pipe(prompt=prompt, init_image=image, strength=0.75).images[0]
    
    return edited_image

# Streamlit app
def main():
    st.title("Image Inpainting with Stable Diffusion")

    # Upload an image
    uploaded_file = st.file_uploader("Upload an image", type=["png", "jpg", "jpeg"])

    # Input prompt for the image editing
    prompt = st.text_input("Enter your prompt", "Change the color of the dragon and add fire in the background")

    if uploaded_file is not None and prompt:
        # Open the uploaded image
        image = Image.open(uploaded_file)

        # Process the image based on the prompt
        with st.spinner("Processing the image..."):
            edited_image = process_image(image, prompt)

        # Display the edited image
        st.image(edited_image, caption="Edited Image", use_container_width=True)

if __name__ == "__main__":
    main()