Browse Source

Add error handling for model downloads and reorder imports in session files

Daniel Gatis 3 months ago
parent
commit
2cc8750160
3 changed files with 8 additions and 8 deletions
  1. 4 1
      rembg/bg.py
  2. 2 5
      rembg/sessions/ben_custom.py
  3. 2 2
      rembg/sessions/dis_custom.py

+ 4 - 1
rembg/bg.py

@@ -210,7 +210,10 @@ def download_models(models: tuple[str, ...]) -> None:
             sys.exit(1)
             sys.exit(1)
         else:
         else:
             print(f"Downloading model: {model}")
             print(f"Downloading model: {model}")
-            session.download_models()
+            try:
+                session.download_models()
+            except Exception as e:
+                print(f"Error downloading model: {e}")
 
 
 
 
 def remove(
 def remove(

+ 2 - 5
rembg/sessions/ben_custom.py

@@ -1,16 +1,13 @@
 import os
 import os
 from typing import List
 from typing import List
-import onnxruntime as ort
-import numpy as np
 
 
+import numpy as np
+import onnxruntime as ort
 from PIL import Image
 from PIL import Image
 from PIL.Image import Image as PILImage
 from PIL.Image import Image as PILImage
 
 
 from .base import BaseSession
 from .base import BaseSession
 
 
-import numpy as np
-from PIL import Image
-
 
 
 class BenCustomSession(BaseSession):
 class BenCustomSession(BaseSession):
     """This is a class representing a custom session for the Ben model."""
     """This is a class representing a custom session for the Ben model."""

+ 2 - 2
rembg/sessions/dis_custom.py

@@ -1,8 +1,8 @@
 import os
 import os
 from typing import List
 from typing import List
-import onnxruntime as ort
-import numpy as np
 
 
+import numpy as np
+import onnxruntime as ort
 from PIL import Image
 from PIL import Image
 from PIL.Image import Image as PILImage
 from PIL.Image import Image as PILImage