From aa6fb76e0f25461d5ec3c9ff6f5f3a0751aeee11 Mon Sep 17 00:00:00 2001 From: catscarlet Date: Fri, 28 Feb 2025 17:06:42 +0800 Subject: [PATCH 1/3] ensure CUDAExecutionProvider is called in session when using nvidia gpu. --- rembg/sessions/base.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/rembg/sessions/base.py b/rembg/sessions/base.py index bfcb115..0d97d44 100644 --- a/rembg/sessions/base.py +++ b/rembg/sessions/base.py @@ -13,9 +13,17 @@ class BaseSession: def __init__(self, model_name: str, sess_opts: ort.SessionOptions, *args, **kwargs): """Initialize an instance of the BaseSession class.""" self.model_name = model_name + + device_type = ort.get_device() + if device_type == 'GPU' and 'CUDAExecutionProvider' in ort.get_available_providers(): + providers = ['CUDAExecutionProvider', 'CPUExecutionProvider'] + else: + providers = ['CPUExecutionProvider'] + self.inner_session = ort.InferenceSession( str(self.__class__.download_models(*args, **kwargs)), sess_options=sess_opts, + providers=providers, ) def normalize( From adfe3a06c63daf02b7a3721364e244a898e8e7d0 Mon Sep 17 00:00:00 2001 From: catscarlet Date: Fri, 28 Feb 2025 17:08:34 +0800 Subject: [PATCH 2/3] Update info about gpu acceleration issue. --- README.md | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index ca4ca63..c789efc 100644 --- a/README.md +++ b/README.md @@ -100,14 +100,14 @@ pip install "rembg[cli]" # for library + cli Otherwise, install `rembg` with explicit CPU/GPU support. -CPU support: +### CPU support: ```bash pip install rembg[cpu] # for library pip install "rembg[cpu,cli]" # for library + cli ``` -GPU support: +### GPU support: First of all, you need to check if your system supports the `onnxruntime-gpu`. @@ -124,6 +124,8 @@ pip install "rembg[gpu]" # for library pip install "rembg[gpu,cli]" # for library + cli ``` +Nvidia GPU may require onnxruntime-gpu, cuda, and cudnn-devel. [#668](https://github.com/danielgatis/rembg/issues/668#issuecomment-2689830314) . If rembg[gpu] couldn't work probably and your can't install cuda or cudnn-devel, use rembg[cpu] and onnxruntime instead. + ## Usage as a cli After the installation step you can use rembg just typing `rembg` in your terminal window. From 1b37f6d3a4cb995e564051cccbba2ca430a967b4 Mon Sep 17 00:00:00 2001 From: catscarlet Date: Fri, 28 Feb 2025 17:12:24 +0800 Subject: [PATCH 3/3] Add info about docker support --- README.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/README.md b/README.md index c789efc..e38e400 100644 --- a/README.md +++ b/README.md @@ -348,6 +348,8 @@ Try this: docker run -v path/to/input:/rembg danielgatis/rembg i input.png path/to/output/output.png ``` +Notice: Right now docker version only support CPU Acceleration. + ## Models All models are downloaded and saved in the user home folder in the `.u2net` directory.