app.py

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")

# Function to process the image with the provided prompt
def process_image(image, prompt):
    # Ensure the image is in the correct format (PIL.Image)
    if isinstance(image, torch.Tensor):
        image = Image.fromarray(image.numpy())
    elif isinstance(image, np.ndarray):
        image = Image.fromarray(image)
    elif not isinstance(image, Image.Image):
        raise ValueError("The image should be either a PIL Image or a numpy array.")

    # Resize image if needed (e.g., to 512x512 for compatibility with the model)
    image = image.resize((512, 512))

    # Preprocess the image
    image_tensor = pipe.image_processor.preprocess(image, return_tensors="pt").pixel_values

    # Use the model pipeline for inpainting
    edited_image = pipe(prompt=prompt, init_image=image_tensor, 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_column_width=True)

if __name__ == "__main__":
    main()