add erase tab

app.py

@@ -8,6 +8,7 @@ from diffusers import FluxFillPipeline
 from PIL import Image, ImageOps
 from sam2.sam2_image_predictor import SAM2ImagePredictor
 import numpy as np
+from simple_lama_inpainting import SimpleLama
 
 torch.set_float32_matmul_precision(["high", "highest"][0])
 
@@ -29,6 +30,7 @@ pipe = FluxFillPipeline.from_pretrained(
     "black-forest-labs/FLUX.1-Fill-dev", torch_dtype=torch.bfloat16
 ).to("cuda")
 
+simple_lama = SimpleLama()
 
 def prepare_image_and_mask(
     image,
@@ -147,6 +149,10 @@ def mask_generation(image=None, d=None):
 
     return out
 
+def erase(image=None,mask=None) : 
+    image = load_img(image)
+    mask = load_img(mask).convert("L")
+    return simple_lama(image, mask)
 
 @spaces.GPU
 def main(*args):
@@ -160,6 +166,8 @@ def main(*args):
         return inpaint(*args)
     elif api_num == 4:
         return mask_generation(*args)
+    elif api_num == 5 : 
+        return erase(*args)
 
 
 rmbg_tab = gr.Interface(
@@ -233,6 +241,25 @@ sam2_tab = gr.Interface(
     cache_examples=False,
 )
 
+erase_tab = gr.Interface(
+    main,
+    inputs=[
+        gr.Number(5, interactive=False),
+        gr.Image(type="pil"),
+        gr.Image(type="pil"),
+    ],
+    outputs=gr.Image(),
+    examples=[
+        [
+            5,
+            "./assets/rocket.png",
+            "./assets/Inpainting mask.png",
+        ]
+    ],
+    api_name="erase",
+    cache_examples=False,
+)
+
 demo = gr.TabbedInterface(
     [rmbg_tab, outpaint_tab, inpaint_tab, sam2_tab],
     ["remove background", "outpainting", "inpainting", "sam2"],

requirements.txt

@@ -20,4 +20,5 @@ huggingface_hub
 sentencepiece
 einops
 git+https://github.com/facebookresearch/sam2.git
-matplotlib
+matplotlib
+simple-lama-inpainting