app.py

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

# 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 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 to the required size (512x512)
    image = image.resize((512, 512))

    # Convert the PIL image to the format that the model expects
    image = np.array(image)  # Convert to numpy array
    image = torch.from_numpy(image).unsqueeze(0).float()  # Convert to tensor and add batch dimension

    # 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()