Change model precision to float32 for CPU compatibility

app.py

@@ -10,15 +10,25 @@ from SegBody import segment_body  # Import the segmentation function
 # Check if CUDA is available and set the device accordingly
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# Load models with fp32 for CPU compatibility
-vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float32)  # Use float32 for CPU
-pipeline = AutoPipelineForInpainting.from_pretrained(
-    "diffusers/stable-diffusion-xl-1.0-inpainting-0.1", 
-    vae=vae, 
-    torch_dtype=torch.float32,  # Use float32 for CPU
-    variant="fp32",  # Use fp32 for CPU
-    use_safetensors=True
-).to(device)  # Ensure it uses the appropriate device (CPU or GPU)
+# Load models with correct precision for CPU or GPU
+if device == "cuda":
+    vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)  # Use fp16 for GPU
+    pipeline = AutoPipelineForInpainting.from_pretrained(
+        "diffusers/stable-diffusion-xl-1.0-inpainting-0.1", 
+        vae=vae, 
+        torch_dtype=torch.float16,  # Use fp16 for GPU
+        variant="fp16",  # Correct variant for GPU
+        use_safetensors=True
+    ).to(device)  # Ensure it uses the GPU
+else:
+    vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float32)  # Use fp32 for CPU
+    pipeline = AutoPipelineForInpainting.from_pretrained(
+        "diffusers/stable-diffusion-xl-1.0-inpainting-0.1", 
+        vae=vae, 
+        torch_dtype=torch.float32,  # Use fp32 for CPU
+        variant="fp32",  # Use fp32 for CPU
+        use_safetensors=True
+    ).to(device)  # Ensure it uses the CPU if no GPU
 
 # Define the inference function
 def inpaint(person_image, garment_image, prompt):