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

fix docker

Browse Source
This commit is contained in:
Daniel Gatis 2023-06-29 21:02:02 -03:00
parent c0b08f831b
commit ccaa9005af
2 changed files with 6 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
rembg/sessions/base.py
View File

@ -28,6 +28,11 @@ class BaseSession:
else: else:
self.providers.extend(_providers) self.providers.extend(_providers)
model_path = kwargs.get("model_path")
if model_path is None:
raise ValueError("model_path is required")
self.inner_session = ort.InferenceSession( self.inner_session = ort.InferenceSession(
str(self.__class__.download_models(*args, **kwargs)), str(self.__class__.download_models(*args, **kwargs)),
providers=self.providers, providers=self.providers,

3
rembg/sessions/u2net_custom.py
View File

@ -34,9 +34,8 @@ class U2netCustomSession(BaseSession):
@classmethod @classmethod
def download_models(cls, *args, **kwargs): def download_models(cls, *args, **kwargs):
model_path = kwargs.get("model_path") model_path = kwargs.get("model_path")
if model_path is None: if model_path is None:
raise ValueError("model_path is required") return
return os.path.abspath(os.path.expanduser(model_path)) return os.path.abspath(os.path.expanduser(model_path))