bug fix

app.py

@@ -32,23 +32,15 @@ clipaway = CLIPAway(
 )
 
 def dilate_mask(mask, kernel_size=5, iterations=5):
-    mask = mask.convert("L")
+    mask = mask.convert("L").resize((512, 512), Image.NEAREST)
     kernel = np.ones((kernel_size, kernel_size), np.uint8)
     mask = cv2.dilate(np.array(mask), kernel, iterations=iterations)
     return Image.fromarray(mask)
 
-def combine_masks(uploaded_mask, sketched_mask):
-    if uploaded_mask is not None:
-        return uploaded_mask
-    elif sketched_mask is not None:
-        return sketched_mask
-    else:
-        raise ValueError("Please provide a mask")
-
 @spaces.GPU
 def remove_obj(image, uploaded_mask, seed):
-    image_pil, sketched_mask = image["image"], image["mask"]
-    mask = dilate_mask(combine_masks(uploaded_mask, sketched_mask))
+    image_pil = image["image"].resize((512, 512), Image.ANTIALIAS)
+    mask = dilate_mask(uploaded_mask)
     seed = int(seed)
     latents = torch.randn((1, 4, 64, 64), generator=torch.Generator().manual_seed(seed)).to("cuda")
     final_image = clipaway.generate(
@@ -62,22 +54,21 @@ examples = [
     ["gradio_examples/images/1.jpg", "gradio_examples/masks/1.png", 42],
     ["gradio_examples/images/2.jpg", "gradio_examples/masks/2.png", 42],
     ["gradio_examples/images/3.jpg", "gradio_examples/masks/3.png", 464],
-    ["gradio_examples/images/4.jpg", "gradio_examples/masks/4.png", 2024],
 ]
 
-with gr.Blocks() as demo:
+with gr.Blocks(theme="gradio/monochrome") as demo:
     gr.Markdown("<h1 style='text-align:center'>CLIPAway: Harmonizing Focused Embeddings for Removing Objects via Diffusion Models</h1>")
     gr.Markdown("""
         <div style='display:flex; justify-content:center; align-items:center;'>
-            <a href='https://arxiv.org/abs/2406.09368' style="margin-right:10px; color:white;">Paper</a> |
-            <a href='https://yigitekin.github.io/CLIPAway/' style="margin:10px; color:white;">Project Website</a> |
-            <a href='https://github.com/YigitEkin/CLIPAway' style="margin-left:10px; color:white;">GitHub</a>
+            <a href='https://arxiv.org/abs/2406.09368' style="margin-right:10px;">Paper</a> |
+            <a href='https://yigitekin.github.io/CLIPAway/' style="margin:10px;">Project Website</a> |
+            <a href='https://github.com/YigitEkin/CLIPAway' style="margin-left:10px;">GitHub</a>
         </div>
     """)
     gr.Markdown("""
             This application allows you to remove objects from images using the CLIPAway method with diffusion models.
             To use this tool:
-            1. Upload an image.
+            1. Upload an image. (NOTE: We expect a 512x512 image, if you upload a different size, it will be resized to 512x512 which can affect the results.)
             2. Upload a pre-defined mask if you have one. (If you don't have a mask, and want to sketch one, 
             we have provided a gradio demo in our github repository. <br/> Unfortunately, we cannot provide it here due to the compatibility issues with zerogpu.)
             3. Set the seed for reproducibility (default is 42).