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)
         else:
             print(f"Downloading model: {model}")
-            session.download_models()
+            try:
+                session.download_models()
+            except Exception as e:
+                print(f"Error downloading model: {e}")
 
 
 def remove(

+ 2 - 5
rembg/sessions/ben_custom.py

@@ -1,16 +1,13 @@
 import os
 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.Image import Image as PILImage
 
 from .base import BaseSession
 
-import numpy as np
-from PIL import Image
-
 
 class BenCustomSession(BaseSession):
     """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
 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.Image import Image as PILImage