Update app.py

app.py

@@ -31,15 +31,21 @@ def dilate_mask(mask, kernel_size=5, iterations=5):
     mask = cv2.dilate(np.array(mask), kernel, iterations=iterations)
     return Image.fromarray(mask)
 
-@spaces.GPU
-def remove_obj(image, uploaded_mask, seed):
-    image_pil = image.resize((512, 512), Image.LANCZOS)
+def remove_obj(image, seed):
+    alpha_channel = img["layers"][0][:, :, 3]
+    mask = np.where(alpha_channel == 0, 0, 255).astype(np.uint8)
+    uploaded_mask = Image.fromarray(mask)
+    background = Image.fromarray(img["background"])
+    
     mask = dilate_mask(uploaded_mask)
     seed = int(seed)
-    latents = torch.randn((1, 4, 64, 64), generator=torch.Generator().manual_seed(seed)).to("cuda")
+    latents = torch.randn((1, 4, 64, 64), generator=torch.Generator().manual_seed(seed)).to("cpu")
     final_image = clipaway.generate(
         prompt=[""], scale=1, seed=seed,
-        pil_image=[image_pil], alpha=[mask], strength=1, latents=latents
+        pil_image=[background], 
+        alpha=[mask], 
+        strength=1, 
+        latents=latents
     )[0]
     return final_image
 
@@ -73,8 +79,7 @@ with gr.Blocks(theme="gradio/monochrome") as demo:
     
     with gr.Row():
         with gr.Column():
-            image_input = gr.Image(label="Upload Image and Sketch Mask", type="pil", tool="sketch")
-            uploaded_mask = gr.Image(label="Upload Mask", type="pil", image_mode="L")
+            image_input = gr.ImageMask(label="Upload Image and Sketch Mask", type="pil")
             seed_input = gr.Number(value=42, label="Seed")
             process_button = gr.Button("Remove Object")
         with gr.Column():
@@ -82,7 +87,7 @@ with gr.Blocks(theme="gradio/monochrome") as demo:
     
     process_button.click(
         fn=remove_obj,
-        inputs=[image_input, uploaded_mask, seed_input],
+        inputs=[image_input, seed_input],
         outputs=result_image
     )