AI/rembg
1
0
Fork 0
mirror of https://git.mirrors.martin98.com/https://github.com/danielgatis/rembg synced 2025-08-18 05:55:53 +08:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #728 from catscarlet/main

Browse Source 
Fix "Always using CPU, ingore GPU" issue, caused by onnxruntime's strange behavor
This commit is contained in:
Daniel Gatis 2025-03-04 22:26:52 -05:00 committed by GitHub
commit 29d8e28c04
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 14 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
README.md
View File

@ -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.
@ -346,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.

8
rembg/sessions/base.py
View File

@ -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(