fix gpu

.dockerignore

@@ -3,6 +3,8 @@
 !setup.py
 !setup.cfg
 !requirements.txt
+!requirements-cpu.txt
+!requirements-gpu.txt
 !versioneer.py
 !README.md
 !pixel.png

Dockerfile

@@ -12,9 +12,7 @@ WORKDIR /rembg
 COPY . .
 
 RUN GPU=1 pip3 install .
-
-# First run to compile AOT & download model
-RUN rembg pixel.png >/dev/null
+RUN rembg i pixel.png > /dev/null
 
 ENTRYPOINT ["rembg"]
 CMD []

rembg/detect.py

@@ -30,7 +30,7 @@ def ort_session(model_name: str) -> ort.InferenceSession:
     with redirect_stdout(sys.stderr):
         gdown.cached_download(url, path, md5=md5)
 
-    return ort.InferenceSession(path)
+    return ort.InferenceSession(path, providers=ort.get_available_providers())
 
 
 def norm_pred(d: np.ndarray) -> np.ndarray:

requirements.txt

@@ -4,7 +4,7 @@ click==8.0.3
 fastapi==0.72.0
 filetype==1.0.9
 gdown==4.2.0
-numpy==1.22.1
+numpy==1.21.5
 pillow==9.0.0
 pymatting==1.1.5
 python-multipart==0.0.5