Change base image to cuda

Dockerfile

@@ -1,9 +1,15 @@
-FROM nvcr.io/nvidia/tensorrt:24.02-py3
+FROM nvidia/cuda:12.3.2-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()'
 

setup.py

@@ -36,7 +36,7 @@
         "twine",
         "wheel",
     ],
-    "gpu": [""],
+    "gpu": ["onnxruntime-gpu"],
     "cli": [
         "aiohttp",
         "asyncer",
@@ -85,7 +85,7 @@
     python_requires=">=3.8, <3.13",
     packages=find_packages(),
     install_requires=install_requires,
-    # dependency_links=['https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/onnxruntime-cuda-12/pypi/simple/'],
+    dependency_links=['https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/onnxruntime-cuda-12/pypi/simple/'],
     entry_points=entry_points,
     extras_require=extras_require,
     version=versioneer.get_version(),