add remove_bg_device

iopaint/plugins/remove_bg.py

@@ -6,7 +6,13 @@ import torch
 from torch.hub import get_dir
 
 from iopaint.plugins.base_plugin import BasePlugin
-from iopaint.schema import RunPluginRequest, RemoveBGModel
+from iopaint.schema import Device, RunPluginRequest, RemoveBGModel
+
+
+def _rmbg_remove(device, *args, **kwargs):
+    from rembg import remove
+
+    return remove(*args, **kwargs)
 
 
 class RemoveBG(BasePlugin):
@@ -14,9 +20,10 @@ class RemoveBG(BasePlugin):
     support_gen_mask = True
     support_gen_image = True
 
-    def __init__(self, model_name):
+    def __init__(self, model_name, device):
         super().__init__()
         self.model_name = model_name
+        self.device = device
 
         if model_name.startswith("birefnet"):
             import rembg
@@ -33,13 +40,15 @@ class RemoveBG(BasePlugin):
         self._init_session(model_name)
 
     def _init_session(self, model_name: str):
+        self.device_warning()
+
         if model_name == RemoveBGModel.briaai_rmbg_1_4:
             from iopaint.plugins.briarmbg import (
                 create_briarmbg_session,
                 briarmbg_process,
             )
 
-            self.session = create_briarmbg_session()
+            self.session = create_briarmbg_session().to(self.device)
             self.remove = briarmbg_process
         elif model_name == RemoveBGModel.briaai_rmbg_2_0:
             from iopaint.plugins.briarmbg2 import (
@@ -47,13 +56,13 @@ class RemoveBG(BasePlugin):
                 briarmbg2_process,
             )
 
-            self.session = create_briarmbg2_session()
+            self.session = create_briarmbg2_session().to(self.device)
             self.remove = briarmbg2_process
         else:
-            from rembg import new_session, remove
+            from rembg import new_session
 
             self.session = new_session(model_name=model_name)
-            self.remove = remove
+            self.remove = _rmbg_remove
 
     def switch_model(self, new_model_name):
         if self.model_name == new_model_name:
@@ -70,7 +79,7 @@ class RemoveBG(BasePlugin):
         bgr_np_img = cv2.cvtColor(rgb_np_img, cv2.COLOR_RGB2BGR)
 
         # return BGRA image
-        output = self.remove(bgr_np_img, session=self.session)
+        output = self.remove(self.device, bgr_np_img, session=self.session)
         return cv2.cvtColor(output, cv2.COLOR_BGRA2RGBA)
 
     @torch.inference_mode()
@@ -78,7 +87,9 @@ class RemoveBG(BasePlugin):
         bgr_np_img = cv2.cvtColor(rgb_np_img, cv2.COLOR_RGB2BGR)
 
         # return BGR image, 255 means foreground, 0 means background
-        output = self.remove(bgr_np_img, session=self.session, only_mask=True)
+        output = self.remove(
+            self.device, bgr_np_img, session=self.session, only_mask=True
+        )
         return output
 
     def check_dep(self):
@@ -86,3 +97,12 @@ class RemoveBG(BasePlugin):
             import rembg
         except ImportError:
             return "RemoveBG is not installed, please install it first. pip install -U rembg"
+
+    def device_warning(self):
+        if self.device == Device.cuda and self.model_name not in [
+            RemoveBGModel.briaai_rmbg_1_4,
+            RemoveBGModel.briaai_rmbg_2_0,
+        ]:
+            logger.warning(
+                f"remove_bg_device=cuda only supports briaai models({RemoveBGModel.briaai_rmbg_1_4.value}/{RemoveBGModel.briaai_rmbg_2_0.value})"
+            )