test add GPU

Dockerfile

@@ -1,4 +1,4 @@
-FROM ubuntu:22.04
+FROM nvidia/cuda:12.1.1-cudnn8-devel-ubuntu22.04
 
 ENV DEBIAN_FRONTEND=noninteractive
 RUN apt-get update && \
@@ -23,7 +23,8 @@ RUN apt-get update && \
     libxmlsec1-dev \
     libffi-dev \
     liblzma-dev \
-    ffmpeg
+    ffmpeg \
+    nvidia-driver-515
 
 RUN useradd -m -u 1000 user
 USER user
@@ -39,7 +40,7 @@ RUN pyenv install ${PYTHON_VERSION} && \
     pyenv rehash && \
     pip install --no-cache-dir -U pip setuptools wheel && \
     pip install "huggingface-hub" "hf-transfer" "gradio[oauth]>=4.28.0" "gradio_huggingfacehub_search==0.0.7" "APScheduler" && \
-    pip install --pre -U -f https://mlc.ai/wheels mlc-llm-nightly-cpu mlc-ai-nightly-cpu
+    python -m pip install --pre -U -f https://mlc.ai/wheels mlc-llm-nightly-cu121 mlc-ai-nightly-cu121
     
 COPY --chown=1000 . ${HOME}/app
 
@@ -52,6 +53,8 @@ ENV PYTHONPATH=${HOME}/app \
     GRADIO_THEME=huggingface \
     TQDM_POSITION=-1 \
     TQDM_MININTERVAL=1 \
-    SYSTEM=spaces
+    SYSTEM=spaces \
+    LD_LIBRARY_PATH=/usr/local/cuda/lib64:${LD_LIBRARY_PATH} \
+    PATH=/usr/local/nvidia/bin:${PATH}
 
 ENTRYPOINT /bin/sh start.sh