2
0
Daniel Gatis 2 жил өмнө
parent
commit
ccaa9005af

+ 5 - 0
rembg/sessions/base.py

@@ -28,6 +28,11 @@ class BaseSession:
         else:
             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(
             str(self.__class__.download_models(*args, **kwargs)),
             providers=self.providers,

+ 1 - 2
rembg/sessions/u2net_custom.py

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