Updated App.py

app.py

@@ -3,20 +3,27 @@ from PIL import Image
 from diffusers import AutoPipelineForInpainting, AutoencoderKL
 import torch
 
+# Check if CUDA is available and set the device accordingly
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
 # Load models
-vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
-pipeline = AutoPipelineForInpainting.from_pretrained("diffusers/stable-diffusion-xl-1.0-inpainting-0.1", 
-                                                     vae=vae, 
-                                                     torch_dtype=torch.float16, 
-                                                     variant="fp16", 
-                                                     use_safetensors=True).to("cuda")
+vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float32)  # Use float32 for CPU compatibility
+pipeline = AutoPipelineForInpainting.from_pretrained(
+    "diffusers/stable-diffusion-xl-1.0-inpainting-0.1", 
+    vae=vae, 
+    torch_dtype=torch.float32,  # Use float32 for CPU compatibility
+    variant="fp32",  # Use fp32 for CPU
+    use_safetensors=True
+).to(device)  # Ensure it uses the appropriate device (CPU or GPU)
 
 # Define the inference function
 def inpaint(prompt, image, mask_image, ip_image):
+    # Preprocess the images by resizing them to 512x512
     image = image.convert("RGB").resize((512, 512))
     mask_image = mask_image.convert("RGB").resize((512, 512))
     ip_image = ip_image.convert("RGB").resize((512, 512))
 
+    # Perform inpainting using the pipeline
     results = pipeline(
         prompt=prompt,
         negative_prompt="ugly, bad quality, bad anatomy",
@@ -27,6 +34,7 @@ def inpaint(prompt, image, mask_image, ip_image):
         guidance_scale=8.0,
         num_inference_steps=100
     )
+
     return results.images[0]
 
 # Set up the Gradio interface