kk

app.py

@@ -29,11 +29,12 @@ def load_models():
     
     if sd_pipe is None:
         print("Loading Stable Diffusion model...")
+        device = "cuda" if torch.cuda.is_available() else "cpu"
         sd_pipe = StableDiffusionInpaintPipeline.from_pretrained(
             "stabilityai/stable-diffusion-2-inpainting",
-            torch_dtype=torch.float16,
+            torch_dtype=torch.float16 if device == "cuda" else torch.float32,
             safety_checker=None
-        ).to("cuda")
+        ).to(device)
     
     return predictor, sd_pipe