Update rembg dependency and add BiRefNet models in the list of available RemBG models

.gitignore

@@ -1 +1,2 @@
 __pycache__
+.vscode

README.md

@@ -6,6 +6,10 @@ Extension for [webui](https://github.com/AUTOMATIC1111/stable-diffusion-webui).
 
 Find the UI for rembg in the Extras tab after installing the extension.
 
+You'll be able to select the following models:
+- "Classic" models: inference is fast, result is good.
+- "BiRefNet" models: inference may take more time but result might be better.
+
 # Installation
 
 Install from webui's Extensions tab.

install.py

@@ -1,7 +1,19 @@
 import launch
+from importlib import metadata
+
+
+rembg_expected_version = "2.0.59"
+
+
+try:
+    rembg_installed_version = metadata.version("rembg")
+except Exception:
+    rembg_installed_version = None
+
+
+if rembg_installed_version != rembg_expected_version:
+    launch.run_pip(f"install rembg=={rembg_expected_version} --no-deps", "rembg")
 
-if not launch.is_installed("rembg"):
-    launch.run_pip("install rembg==2.0.50 --no-deps", "rembg")
 
 for dep in ['onnxruntime', 'pymatting', 'pooch']:
     if not launch.is_installed(dep):

scripts/api.py

@@ -7,12 +7,20 @@
 import rembg
 
 # models = [
-#     "None",
 #     "u2net",
 #     "u2netp",
 #     "u2net_human_seg",
 #     "u2net_cloth_seg",
 #     "silueta",
+#     "isnet-general-use",
+#     "isnet-anime",
+#     "birefnet-general",
+#     "birefnet-general-lite",
+#     "birefnet-portrait",
+#     "birefnet-dis",
+#     "birefnet-hrsod",
+#     "birefnet-cod",
+#     "birefnet-massive",
 # ]
 
 
@@ -34,7 +42,7 @@ async def rembg_remove(
 
         image = rembg.remove(
             input_image,
-            session=rembg.new_session(model),
+            session=rembg.new_session(model, providers=['CUDAExecutionProvider', 'CPUExecutionProvider']),
             only_mask=return_mask,
             alpha_matting=alpha_matting,
             alpha_matting_foreground_threshold=alpha_matting_foreground_threshold,

scripts/postprocessing_rembg.py

@@ -1,14 +1,12 @@
 from modules import scripts_postprocessing, ui_components
 import gradio as gr
 
-from modules.ui_components import FormRow
 from modules.paths_internal import models_path
 import rembg
 import os
 
 models = [
     "None",
-    "isnet-general-use",
     "u2net",
     "u2netp",
     "u2net_human_seg",
@@ -18,13 +16,27 @@
     "isnet-anime",
 ]
 
+birefnet_models = [
+    "None",
+    "birefnet-general",
+    "birefnet-general-lite",
+    "birefnet-portrait",
+    "birefnet-dis",
+    "birefnet-hrsod",
+    "birefnet-cod",
+    "birefnet-massive",
+]
+
 class ScriptPostprocessingUpscale(scripts_postprocessing.ScriptPostprocessing):
     name = "Rembg"
     order = 20000
     model = None
 
     def ui(self):
         with ui_components.InputAccordion(False, label="Remove background") as enable:
+            with gr.Row():
+                model_type = gr.Radio(["Classic", "BiRefNet"], value="Classic", label="Model type")
+
             with gr.Row():
                 model = gr.Dropdown(label="Remove background", choices=models, value="None")
                 return_mask = gr.Checkbox(label="Return mask", value=False)
@@ -35,6 +47,12 @@ def ui(self):
                 alpha_matting_foreground_threshold = gr.Slider(label="Foreground threshold", minimum=0, maximum=255, step=1, value=240)
                 alpha_matting_background_threshold = gr.Slider(label="Background threshold", minimum=0, maximum=255, step=1, value=10)
 
+            model_type.change(
+                fn=lambda x: gr.update(value="None", choices=birefnet_models if x == "BiRefNet" else models),
+                inputs=[model_type],
+                outputs=[model],
+            )
+
             alpha_matting.change(
                 fn=lambda x: gr.update(visible=x),
                 inputs=[alpha_matting],
@@ -63,7 +81,7 @@ def process(self, pp: scripts_postprocessing.PostprocessedImage, enable, model,
 
         pp.image = rembg.remove(
             pp.image,
-            session=rembg.new_session(model),
+            session=rembg.new_session(model, providers=['CUDAExecutionProvider', 'CPUExecutionProvider']),
             only_mask=return_mask,
             alpha_matting=alpha_matting,
             alpha_matting_foreground_threshold=alpha_matting_foreground_threshold,