Updated Dockerfile for GPU Dependencies

Dockerfile

@@ -1,20 +1,34 @@
-FROM python:3.10
+# Use an appropriate base image with Ubuntu
+FROM nvidia/cuda:11.2.1-base-ubuntu20.04   # Adjust the CUDA version as needed
 
-RUN apt-get update && apt-get install -y ffmpeg && apt-get install -y yt-dlp
+# Set environment variables for CUDA
+ENV CUDA_HOME=/usr/local/cuda
+ENV PATH=$CUDA_HOME/bin:$PATH
+ENV LD_LIBRARY_PATH=$CUDA_HOME/lib64:$LD_LIBRARY_PATH
 
-WORKDIR /code
+# Install NVIDIA driver
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    nvidia-driver-460.106   # Use the specific driver version you found
 
-COPY ./requirements.txt /code/requirements.txt
+# Additional system dependencies
+RUN apt-get install -y ffmpeg && apt-get install -y yt-dlp
+
+# Switch back to the root user to install Python packages
+USER root
 
+# Copy your requirements.txt and install Python packages
+COPY ./requirements.txt /code/requirements.txt
 RUN pip install --no-cache-dir --upgrade -r /code/requirements.txt
 
+# Create a non-root user and set environment variables
 RUN useradd -m -u 1000 user
 USER user
 ENV HOME=/home/user \
-	PATH=/home/user/.local/bin:$PATH
+    PATH=/home/user/.local/bin:$PATH
 
+# Set the working directory and copy your application files
 WORKDIR $HOME/app
-
 COPY --chown=user . $HOME/app
 
-CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "7860"]
+# Specify the command to run your application
+CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "7860"]