add remove_bg_device

iopaint/plugins/__init__.py

@@ -16,6 +16,7 @@ def build_plugins(
     interactive_seg_model: InteractiveSegModel,
     interactive_seg_device: Device,
     enable_remove_bg: bool,
+    remove_bg_device: Device,
     remove_bg_model: str,
     enable_anime_seg: bool,
     enable_realesrgan: bool,
@@ -36,7 +37,7 @@ def build_plugins(
 
     if enable_remove_bg:
         logger.info(f"Initialize {RemoveBG.name} plugin")
-        plugins[RemoveBG.name] = RemoveBG(remove_bg_model)
+        plugins[RemoveBG.name] = RemoveBG(remove_bg_model, remove_bg_device)
 
     if enable_anime_seg:
         logger.info(f"Initialize {AnimeSeg.name} plugin")

iopaint/plugins/briarmbg.py

@@ -480,7 +480,7 @@ def create_briarmbg_session():
     return net
 
 
-def briarmbg_process(bgr_np_image, session, only_mask=False):
+def briarmbg_process(device, bgr_np_image, session, only_mask=False):
     # prepare input
     orig_bgr_image = Image.fromarray(bgr_np_image)
     w, h = orig_im_size = orig_bgr_image.size
@@ -490,6 +490,7 @@ def briarmbg_process(bgr_np_image, session, only_mask=False):
     im_tensor = torch.unsqueeze(im_tensor, 0)
     im_tensor = torch.divide(im_tensor, 255.0)
     im_tensor = normalize(im_tensor, [0.5, 0.5, 0.5], [1.0, 1.0, 1.0])
+    im_tensor = im_tensor.to(device)
     # inference
     result = session(im_tensor)
     # post process

iopaint/plugins/briarmbg2.py

@@ -10,7 +10,7 @@ def create_briarmbg2_session():
     return birefnet
 
 
-def briarmbg2_process(bgr_np_image, session, only_mask=False):
+def briarmbg2_process(device, bgr_np_image, session, only_mask=False):
     from torchvision import transforms
     from PIL import Image
 
@@ -25,6 +25,7 @@ def briarmbg2_process(bgr_np_image, session, only_mask=False):
     image = Image.fromarray(bgr_np_image)
     image_size = image.size
     input_images = transform_image(image).unsqueeze(0)
+    input_images = input_images.to(device)
 
     # Prediction
     preds = session(input_images)[-1].sigmoid().cpu()

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})"
+            )