update

app.py

@@ -109,6 +109,11 @@ with gr.Blocks(
                         value=2048,
                         step=256,
                     )
+                    rmbg_tag = gr.Radio(
+                        choices=["rembg", "rmbg14"],
+                        value="rembg",
+                        label="Background Removal Model",
+                    )
                 with gr.Row():
                     randomize_seed = gr.Checkbox(
                         label="Randomize Seed", value=False
@@ -207,10 +212,13 @@ with gr.Blocks(
             with gr.Row() as single_image_example:
                 examples = gr.Examples(
                     label="Image Gallery",
-                    examples=sorted(
-                        glob("assets/example_image/*")
-                    ),
-                    inputs=[image_prompt],
+                    examples=[
+                        [image_path]
+                        for image_path in sorted(
+                            glob("assets/example_image/*")
+                        )
+                    ],
+                    inputs=[image_prompt, rmbg_tag],
                     fn=preprocess_image_fn,
                     outputs=[image_prompt, raw_image_cache],
                     run_on_click=True,
@@ -220,9 +228,12 @@ with gr.Blocks(
             with gr.Row(visible=False) as single_sam_image_example:
                 examples = gr.Examples(
                     label="Image Gallery",
-                    examples=sorted(
-                        glob("assets/example_image/*")
-                    ),
+                    examples=[
+                        [image_path]
+                        for image_path in sorted(
+                            glob("assets/example_image/*")
+                        )
+                    ],
                     inputs=[image_prompt_sam],
                     fn=preprocess_sam_image_fn,
                     outputs=[image_prompt_sam, raw_image_cache],
@@ -274,7 +285,7 @@ with gr.Blocks(
 
     image_prompt.upload(
         preprocess_image_fn,
-        inputs=[image_prompt],
+        inputs=[image_prompt, rmbg_tag],
         outputs=[image_prompt, raw_image_cache],
     )
     image_prompt.change(

common.py

@@ -150,8 +150,8 @@ def download_kolors_weights() -> None:
 
 
 if os.getenv("GRADIO_APP") == "imageto3d":
-    # RBG_REMOVER = RembgRemover()
-    RBG_REMOVER = BMGG14Remover()
+    RBG_REMOVER = RembgRemover()
+    RBG14_REMOVER = BMGG14Remover()
     SAM_PREDICTOR = SAMPredictor(model_type="vit_h", device="cpu")
     PIPELINE = TrellisImageTo3DPipeline.from_pretrained(
         "JeffreyXiang/TRELLIS-image-large"
@@ -165,8 +165,8 @@ if os.getenv("GRADIO_APP") == "imageto3d":
         os.path.dirname(os.path.abspath(__file__)), "sessions/imageto3d"
     )
 elif os.getenv("GRADIO_APP") == "textto3d":
-    # RBG_REMOVER = RembgRemover()
-    RBG_REMOVER = BMGG14Remover()
+    RBG_REMOVER = RembgRemover()
+    RBG14_REMOVER = BMGG14Remover()
     PIPELINE = TrellisImageTo3DPipeline.from_pretrained(
         "JeffreyXiang/TRELLIS-image-large"
     )
@@ -289,7 +289,7 @@ def render_video(
 
 @spaces.GPU
 def preprocess_image_fn(
-    image: str | np.ndarray | Image.Image,
+    image: str | np.ndarray | Image.Image, rmbg_tag: str = "rembg"
 ) -> tuple[Image.Image, Image.Image]:
     if isinstance(image, str):
         image = Image.open(image)
@@ -298,7 +298,8 @@ def preprocess_image_fn(
 
     image_cache = image.copy().resize((512, 512))
 
-    image = RBG_REMOVER(image)
+    bg_remover = RBG_REMOVER if rmbg_tag == "rembg" else RBG14_REMOVER
+    image = bg_remover(image)
     image = trellis_preprocess(image)
 
     return image, image_cache
@@ -710,6 +711,7 @@ def text2image_fn(
     ip_image: Image.Image | str = None,
     ip_adapt_scale: float = 0.3,
     image_wh: int | tuple[int, int] = [1024, 1024],
+    rmbg_tag: str = "rembg",
     n_sample: int = 3,
     req: gr.Request = None,
 ):
@@ -736,7 +738,7 @@ def text2image_fn(
 
     for idx in range(len(images)):
         image = images[idx]
-        images[idx], _ = preprocess_image_fn(image)
+        images[idx], _ = preprocess_image_fn(image, rmbg_tag)
 
     save_paths = []
     for idx, image in enumerate(images):