AI/rembg
1
0
Fork 0
mirror of https://git.mirrors.martin98.com/https://github.com/danielgatis/rembg synced 2025-08-19 06:35:57 +08:00
Code Issues Packages Projects Releases Wiki Activity

ensure CUDAExecutionProvider is called in session when using nvidia gpu.

Browse Source
This commit is contained in:
catscarlet 2025-02-28 17:06:42 +08:00
parent 9079508935
commit aa6fb76e0f
1 changed files with 8 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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): def __init__(self, model_name: str, sess_opts: ort.SessionOptions, *args, **kwargs):
"""Initialize an instance of the BaseSession class.""" """Initialize an instance of the BaseSession class."""
self.model_name = model_name 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( self.inner_session = ort.InferenceSession(
str(self.__class__.download_models(*args, **kwargs)), str(self.__class__.download_models(*args, **kwargs)),
sess_options=sess_opts, sess_options=sess_opts,
providers=providers,
) )
def normalize( def normalize(