Update app.py

app.py

@@ -3,8 +3,6 @@ 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")
@@ -12,21 +10,15 @@ pipe.to("cuda" if torch.cuda.is_available() else "cpu")  # Move model to GPU if
 
 # 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.")
-
+    # Ensure the image is a PIL Image
+    if isinstance(image, np.ndarray):
+        image = Image.fromarray(image)  # Convert numpy array to PIL Image
+    elif isinstance(image, torch.Tensor):
+        image = Image.fromarray(image.numpy())  # Convert torch tensor to PIL Image
+    
     # 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]