Update app.py

app.py

@@ -17,7 +17,7 @@ from compel import Compel
 from diffusers import EulerDiscreteScheduler
 
 # Device configuration
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+device = torch.device("cpu")  # Ensure everything is set to run on CPU
 
 # Constants for colors
 BG_COLOR = (0, 0, 0, 255)  # gray with full opacity
@@ -31,14 +31,14 @@ options = vision.ImageSegmenterOptions(base_options=base_options,
 # Initialize ControlNet inpainting pipeline
 controlnet = ControlNetModel.from_pretrained(
     'lllyasviel/control_v11p_sd15_inpaint',
-    torch_dtype=torch.float16,
+    torch_dtype=torch.float32,  # Use float32 for CPU
 ).to(device)
 
 pipe = StableDiffusionControlNetInpaintPipeline.from_pretrained(
     'runwayml/stable-diffusion-v1-5',
     safety_checker=None,
     controlnet=controlnet,
-    torch_dtype=torch.float16,
+    torch_dtype=torch.float32,  # Use float32 for CPU
 ).to(device)
 
 # Set the K_EULER scheduler
@@ -79,7 +79,6 @@ def resize_image(image, max_size=1536):
         image = cv2.resize(image, new_size, interpolation=cv2.INTER_AREA)
     return image
 
-@spaces.GPU(duration=60)
 # Function to inpaint the hair area using ControlNet
 def inpaint_hair(image, prompt):
     # Only resize the input image if it's larger than 1536 in any dimension
@@ -101,20 +100,19 @@ def inpaint_hair(image, prompt):
     inpaint_condition = torch.from_numpy(np.expand_dims(image_np, 0).transpose(0, 3, 1, 2)).to(device)
 
     # Generate inpainted image
-    generator = torch.Generator(device).manual_seed(42)
+    generator = torch.manual_seed(42)
     negative_prompt = "lowres, bad quality, poor quality"
     
-    with torch.cuda.amp.autocast():  # Enable automatic mixed precision
-        output = pipe(
-            prompt=prompt,
-            negative_prompt=negative_prompt,
-            image=image_pil,
-            mask_image=mask_pil,
-            control_image=inpaint_condition,
-            num_inference_steps=25,
-            guidance_scale=7.5,
-            generator=generator
-        ).images[0]
+    output = pipe(
+        prompt=prompt,
+        negative_prompt=negative_prompt,
+        image=image_pil,
+        mask_image=mask_pil,
+        control_image=inpaint_condition,
+        num_inference_steps=25,
+        guidance_scale=7.5,
+        generator=generator
+    ).images[0]
 
     return np.array(output)
 
@@ -132,4 +130,4 @@ iface = gr.Interface(
 )
 
 if __name__ == "__main__":
-    iface.launch()
+    iface.launch()