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

# Load and display an image in Streamlit
def load_image(image_path):
    # Open image
    image = Image.open(image_path).convert('RGB')
    return image

# Main function to process the image
def process_image(image, prompt):
    # Create the pipeline using Stable Diffusion
    pipe = StableDiffusionInpaintPipeline.from_pretrained("stabilityai/stable-diffusion-2-inpainting")
    
    # If using GPU, send the pipeline to CUDA
    pipe.to("cuda" if torch.cuda.is_available() else "cpu")

    # Perform inpainting (change color and add fire in the background)
    edited_image = pipe(prompt=prompt, init_image=image, strength=0.75).images[0]
    
    return edited_image

# Streamlit Interface
def main():
    # Upload image
    uploaded_image = st.file_uploader("Upload an image", type=["jpg", "png", "jpeg"])
    
    if uploaded_image is not None:
        # Load image
        image = load_image(uploaded_image)

        # Display original image
        st.image(image, caption="Original Image", use_container_width=True)

        # Define the prompt
        prompt = "change the color of dragon and add fire in the background"
        
        # Process the image based on the prompt
        edited_image = process_image(image, prompt)

        # Display the edited image
        st.image(edited_image, caption="Edited Image", use_container_width=True)
    else:
        st.write("Please upload an image to begin.")

if __name__ == "__main__":
    main()