mirror of
https://git.mirrors.martin98.com/https://github.com/danielgatis/rembg
synced 2025-08-18 06:55:53 +08:00
ensure CUDAExecutionProvider is called in session when using nvidia gpu.
This commit is contained in:
parent
9079508935
commit
aa6fb76e0f
@ -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(
|
||||
|
Loading…
x
Reference in New Issue
Block a user