Update app.py

app.py

@@ -3,9 +3,12 @@ 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):
@@ -17,14 +20,16 @@ def process_image(image, prompt):
     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)
+    # Resize image to the required size (512x512)
     image = image.resize((512, 512))
 
-    # Preprocess the image
-    image_tensor = pipe.image_processor.preprocess(image, return_tensors="pt").pixel_values
+    # 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
 
-    # Use the model pipeline for inpainting
-    edited_image = pipe(prompt=prompt, init_image=image_tensor, strength=0.75).images[0]
+    # Process the image through the pipeline
+    edited_image = pipe(prompt=prompt, init_image=image, strength=0.75).images[0]
+    
     return edited_image
 
 # Streamlit app
@@ -46,7 +51,7 @@ def main():
             edited_image = process_image(image, prompt)
 
         # Display the edited image
-        st.image(edited_image, caption="Edited Image", use_column_width=True)
+        st.image(edited_image, caption="Edited Image", use_container_width=True)
 
 if __name__ == "__main__":
     main()