Update Dockerfile

Dockerfile

@@ -1,12 +1,13 @@
-# Base image with CUDA support
-FROM nvidia/cuda:12.1.0-cudnn8-runtime-ubuntu22.04
+FROM nvidia/cuda:12.1.0-cudnn8-devel-ubuntu22.04
 
 # Set environment variables
 ENV PYTHONUNBUFFERED=1 \
     DEBIAN_FRONTEND=noninteractive \
     CUDA_HOME=/usr/local/cuda \
     PATH=/usr/local/cuda/bin:$PATH \
-    LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH
+    LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH \
+    NVIDIA_VISIBLE_DEVICES=all \
+    NVIDIA_DRIVER_CAPABILITIES=compute,utility
 
 # Install system dependencies
 RUN apt-get update && apt-get install -y --no-install-recommends \
@@ -14,41 +15,29 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
     python3-pip \
     python3-dev \
     build-essential \
-    curl \
     ffmpeg \
     libsndfile1 \
+    curl \
     && rm -rf /var/lib/apt/lists/*
 
 # Upgrade pip and install build tools
-RUN python3 -m pip install --no-cache-dir --upgrade pip setuptools wheel
-
-# Install PyTorch separately first
-RUN pip3 install --no-cache-dir \
-    torch==2.4.1 \
-    torchaudio==2.4.1 \
-    torchvision==0.19.1 \
-    --extra-index-url https://download.pytorch.org/whl/cu121
+RUN python3 -m pip install --upgrade pip setuptools wheel
 
-# Set up working directory
 WORKDIR /app
 
-# Copy requirements first to leverage Docker cache
 COPY requirements.txt .
 
-# Modify requirements to remove PyTorch packages
-RUN sed -i '/^torch==/d; /^torchaudio==/d; /^torchvision==/d; /^torchdata==/d' requirements.txt
+# Install PyTorch with CUDA support
+RUN pip3 install --no-cache-dir torch==2.1.2 torchvision==0.16.2 torchaudio==2.1.2 --index-url https://download.pytorch.org/whl/cu121
 
-# Install remaining Python dependencies
+# Install other requirements
 RUN pip3 install --no-cache-dir -r requirements.txt
 
-# Copy the rest of your application
 COPY . .
 
-# Verify GPU setup
-RUN python3 -c "import torch; assert torch.cuda.is_available(), 'CUDA not available'"
+# Simple CUDA test
+RUN python3 -c "import torch; print('CUDA available:', torch.cuda.is_available()); print('CUDA version:', torch.version.cuda); print('Device count:', torch.cuda.device_count())"
 
-# Expose port
 EXPOSE 8000
 
-# Start the application
 CMD ["python3", "server.py"]