Change base image to cuda

PPCBee · Mar 21, 2024 · fa11f6d · fa11f6d
1 parent 816302b
commit fa11f6d
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 3 deletions.
diff --git a/Dockerfile b/Dockerfile
@@ -1,11 +1,17 @@
-FROM nvcr.io/nvidia/tensorrt:24.02-py3
+FROM nvidia/cuda:11.8.0-cudnn8-devel-ubuntu22.04
 
 WORKDIR /rembg
 
 COPY . .
 
+# RUN apt-get update
+# RUN apt install software-properties-common -y
+# RUN add-apt-repository ppa:deadsnakes/ppa -y
+RUN apt-get update
+RUN apt-get install -y python-is-python3 python3 python3-pip
+# RUN pip install onnxruntime-gpu --extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/onnxruntime-cuda-12/pypi/simple/
 RUN python -m pip install ".[gpu,cli]"
-RUN python -c 'from rembg.bg import download_models; download_models()'
+# RUN python -c 'from rembg.bg import download_models; download_models()'
 
 EXPOSE 7000
 ENTRYPOINT ["rembg"]

diff --git a/setup.py b/setup.py
@@ -36,7 +36,7 @@
         "twine",
         "wheel",
     ],
-    "gpu": [""],
+    "gpu": ["onnxruntime-gpu"],
     "cli": [
         "aiohttp",
         "asyncer",