Update app.py

app.py

@@ -1,49 +1,52 @@
 import streamlit as st
 from PIL import Image
-from diffusers import StableDiffusionInpaintPipeline
 import torch
+from diffusers import StableDiffusionInpaintPipeline
+import numpy as np
 
-# Load and display an image in Streamlit
-def load_image(image_path):
-    # Open image
-    image = Image.open(image_path).convert('RGB')
-    return image
+# Load the StableDiffusionInpaintPipeline
+pipe = StableDiffusionInpaintPipeline.from_pretrained("stabilityai/stable-diffusion-2-inpainting")
 
-# Main function to process the image
+# Function to process the image with the provided prompt
 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]
-    
+    # 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 Interface
+# Streamlit app
 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"
-        
+    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
-        edited_image = process_image(image, 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)
-    else:
-        st.write("Please upload an image to begin.")
+        st.image(edited_image, caption="Edited Image", use_column_width=True)
 
 if __name__ == "__main__":
     main()