Update app.py

app.py

@@ -31,21 +31,14 @@ def dilate_mask(mask, kernel_size=5, iterations=5):
     mask = cv2.dilate(np.array(mask), kernel, iterations=iterations)
     return Image.fromarray(mask)
 
-def remove_obj(image, seed):
-    alpha_channel = image["layers"][0][:, :, 3]
-    mask = np.where(alpha_channel == 0, 0, 255).astype(np.uint8)
-    uploaded_mask = Image.fromarray(mask)
-    background = Image.fromarray(image["background"])
-    
-    mask = dilate_mask(uploaded_mask)
+def remove_obj(image, uploaded_mask, seed):
+    image_pil, sketched_mask = image["image"], image["mask"]
+    mask = dilate_mask(combine_masks(uploaded_mask, sketched_mask))
     seed = int(seed)
-    latents = torch.randn((1, 4, 64, 64), generator=torch.Generator().manual_seed(seed)).to("cpu")
+    latents = torch.randn((1, 4, 64, 64), generator=torch.Generator().manual_seed(seed)).to("cuda")
     final_image = clipaway.generate(
         prompt=[""], scale=1, seed=seed,
-        pil_image=[background], 
-        alpha=[mask], 
-        strength=1, 
-        latents=latents
+        pil_image=[image_pil], alpha=[mask], strength=1, latents=latents
     )[0]
     return final_image
 
@@ -79,7 +72,8 @@ with gr.Blocks() as demo:
     
     with gr.Row():
         with gr.Column():
-            image_input = gr.ImageMask(label="Upload Image and Sketch Mask", height=700, layers=False)
+            image_input = gr.Image(label="Upload Image and Sketch Mask", type="pil", tool="sketch")
+            uploaded_mask = gr.Image(label="Upload Mask (Optional)", type="pil", optional=True)
             seed_input = gr.Number(value=42, label="Seed")
             process_button = gr.Button("Remove Object")
         with gr.Column():