Răsfoiți Sursa

Merge pull request #728 from catscarlet/main

Fix "Always using CPU, ingore GPU" issue, caused by onnxruntime's strange behavor
Daniel Gatis 6 luni în urmă
părinte
comite
29d8e28c04
2 a modificat fișierele cu 14 adăugiri și 2 ștergeri
  1. 6 2
      README.md
  2. 8 0
      rembg/sessions/base.py

+ 6 - 2
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.
@@ -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 - 0
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(