Dockerfile: test with pytorch image

Dockerfile

@@ -1,71 +1,43 @@
-# ORI --> cuDNN error: CUDNN_STATUS_SUBLIBRARY_VERSION_MISMATCH
-# FROM nvidia/cuda:12.8.1-cudnn-runtime-ubuntu22.04
-# TESTS UNTIL RIGHT PYTHON VERSION
-# FROM nvidia/cuda:11.6.1-devel-ubuntu20.04
-# FROM nvidia/cuda:12.0.0-cudnn8-devel-ubuntu18.04
-# FROM nvidia/cuda:12.4.1-cudnn-runtime-ubuntu22.04
-# FROM nvidia/cuda:12.1.0-cudnn8-devel-ubuntu22.04
-# FROM nvidia/cuda:12.1.0-cudnn8-runtime-ubuntu22.04
-FROM nvidia/cuda:12.4.1-cudnn-devel-ubuntu22.04
+# ✅ PyTorch 2.2.2 (cu121), Python 3.10, cuDNN 9.5.1+
+FROM pytorch/pytorch:2.2.2-cuda12.1-cudnn8-runtime
 
-ENV DEBIAN_FRONTEND=noninteractive
 ENV PYTHONUNBUFFERED=1
+ENV DEBIAN_FRONTEND=noninteractive
 
 WORKDIR /app
 
-ARG EXTRAS
-ARG HF_PRECACHE_DIR
-ARG HF_TKN_FILE
-
+# Optional: add caching envs (like HF cache) if needed
 ENV HF_HOME="/app/tmp/cache/huggingface"
-ENV HF_HUB_CACHE="/app/tmp/cache/huggingface/hub"
-ENV XDG_CACHE_HOME="/app/tmp/cache/huggingface"
+ENV XDG_CACHE_HOME="/app/tmp/cache"
 ENV LIBROSA_CACHE_DIR="/app/tmp/librosa_cache"
 ENV NUMBA_CACHE_DIR="/app/tmp/numba_cache"
 ENV MPLCONFIGDIR="/app/tmp/matplotlib"
 ENV HF_HUB_ETAG_TIMEOUT="600"
 ENV HF_HUB_DOWNLOAD_TIMEOUT="600"
-ENV LD_LIBRARY_PATH="/usr/local/cuda:/usr/lib/x86_64-linux-gnu:${LD_LIBRARY_PATH:-}"
-
-# Install system dependencies
-#RUN apt-get update && \
-#    apt-get install -y ffmpeg git && \
-#    apt-get clean && \
-#    rm -rf /var/lib/apt/lists/*
 
-# 2) Install system dependencies + Python + pip
+# 🧰 System dependencies
 RUN apt-get update && \
     apt-get install -y --no-install-recommends \
-        python3 \
-        python3-pip \
         ffmpeg \
         git \
-        libportaudio2 && \
+        libportaudio2 \
+        python3-dev \
+        portaudio19-dev && \
+    apt-get clean && \
     rm -rf /var/lib/apt/lists/*
 
-# RUN apt-get install -y --no-install-recommends libportaudio2
-# portaudio19-dev
-
-RUN pip install diart sounddevice
-RUN pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121
+# 📦 Python packages
+# Prefer installing torch packages already bundled, to avoid conflicts
+RUN pip install --upgrade pip && \
+    pip install \
+        diart \
+        sounddevice
 
 COPY . .
 
-RUN mkdir -p $HF_HOME && \
-    mkdir -p $HF_HUB_CACHE && \
-    mkdir -p $XDG_CACHE_HOME && \
-    mkdir -p $LIBROSA_CACHE_DIR && \
-    mkdir -p $NUMBA_CACHE_DIR && \
-    mkdir -p $MPLCONFIGDIR && \
-    chmod 777 $HF_HOME && \
-    chmod 777 $HF_HUB_CACHE && \
-    chmod 777 $XDG_CACHE_HOME && \
-    chmod 777 $LIBROSA_CACHE_DIR && \
-    chmod 777 $NUMBA_CACHE_DIR && \
-    chmod 777 $MPLCONFIGDIR && \
-    chmod -R 777 /app
-
-RUN mkdir -p /.cache && chmod -R 777 /.cache
+# 🧼 Set permissions for Hugging Face cache
+RUN mkdir -p $HF_HOME $XDG_CACHE_HOME $LIBROSA_CACHE_DIR $NUMBA_CACHE_DIR /.cache && \
+    chmod -R 777 /app /root/.cache /.cache
 
 # Install WhisperLiveKit directly, allowing for optional dependencies
 #   Note: For gates modedls, need to add your HF toke. See README.md
@@ -111,11 +83,10 @@ RUN if [ -n "$HF_TKN_FILE" ]; then \
     else \
       echo "No Hugging Face token file specified, skipping token setup"; \
     fi
-    
-# Expose port for the transcription server
+
+# ✅ Expose server port
 EXPOSE 8000
 
+# 🚀 Start the app
 ENTRYPOINT ["whisperlivekit-server", "--host", "0.0.0.0"]
-
-# Default args
-CMD ["--model", "tiny", "--diarization"]
+CMD ["--model", "tiny", "--diarization"]