Update app.py

app.py

@@ -14,6 +14,9 @@ from diffusers import StableDiffusionPipeline, StableDiffusionControlNetInpaintP
 from PIL import Image
 from compel import Compel
 
+# Device configuration
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
 # Constants for colors
 BG_COLOR = (0, 0, 0, 255)  # gray with full opacity
 MASK_COLOR = (255, 255, 255, 255)  # white with full opacity
@@ -27,13 +30,13 @@ options = vision.ImageSegmenterOptions(base_options=base_options,
 controlnet = ControlNetModel.from_pretrained(
     'lllyasviel/control_v11p_sd15_inpaint',
     torch_dtype=torch.float16,
-).to("cuda")
+).to(device)
 
 pipe = StableDiffusionControlNetInpaintPipeline.from_pretrained(
     'runwayml/stable-diffusion-v1-5',
     controlnet=controlnet,
     torch_dtype=torch.float16,
-).to("cuda")
+).to(device)
 
 # Function to segment hair and generate mask
 def segment_hair(image):
@@ -74,7 +77,7 @@ def inpaint_hair(image, prompt):
     image_np = np.array(image_pil).astype(np.float32) / 255.0
     mask_np = np.array(mask_pil.convert("L")).astype(np.float32) / 255.0
     image_np[mask_np > 0.5] = -1.0  # Set as masked pixel
-    inpaint_condition = torch.from_numpy(np.expand_dims(image_np, 0).transpose(0, 3, 1, 2)).to("cuda")
+    inpaint_condition = torch.from_numpy(np.expand_dims(image_np, 0).transpose(0, 3, 1, 2)).to(device)
 
     # Generate inpainted image
     generator = torch.Generator("cuda").manual_seed(42)