From fa11f6d4b48f77ee2b1564c332c14e804cc52e0c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=A1n=20Mo=C5=A1a=C5=A5?= <jmosat@gmail.com>
Date: Thu, 21 Mar 2024 11:40:58 +0100
Subject: [PATCH] Change base image to cuda

---
 Dockerfile | 10 ++++++++--
 setup.py   |  2 +-
 2 files changed, 9 insertions(+), 3 deletions(-)

diff --git a/Dockerfile b/Dockerfile
index 317b9b09..362665b2 100644
--- 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
index 3332664e..22c8a516 100644
--- a/setup.py
+++ b/setup.py
@@ -36,7 +36,7 @@
         "twine",
         "wheel",
     ],
-    "gpu": [""],
+    "gpu": ["onnxruntime-gpu"],
     "cli": [
         "aiohttp",
         "asyncer",