Update Dockerfile

Dockerfile

@@ -1,38 +1,12 @@
+# Use the official PyTorch image as the base
 FROM pytorch/pytorch:2.0.1-cuda11.7-cudnn8-runtime
 
-ENV DEBIAN_FRONTEND=noninteractive
-ENV MKL_THREADING_LAYER=GNU
-
-# Install dependencies
-RUN apt-get update && apt-get install -y \
-    git wget libgl1-mesa-glx libglib2.0-0 ffmpeg libx264-dev pkg-config \
-    build-essential cmake g++ python3.9 python3.9-dev python3.9-venv \
-    software-properties-common wget gnupg2 libopenblas-dev && \
-    rm -rf /var/lib/apt/lists/*
-
-# Ensure Python 3.9 is the default version
-RUN update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.9 1 && \
-    update-alternatives --set python3 /usr/bin/python3.9
-
-# Install pip for Python 3.9
-RUN curl -sS https://bootstrap.pypa.io/get-pip.py | python3.9
-
-# Add NVIDIA package repositories
-RUN wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/cuda-keyring_1.0-1_all.deb && \
-    dpkg -i cuda-keyring_1.0-1_all.deb && \
-    apt-get update
-
-# Install CUDA Toolkit and cuDNN
-RUN apt-get install -y cuda-toolkit-11-7 libcudnn8 libcudnn8-dev
-
-# Create a user
-RUN useradd -m -u 1000 user
-USER user
-
-# Set CUDA paths
-ENV HOME=/home/user \
+# Set environment variables
+ENV DEBIAN_FRONTEND=noninteractive \
+    MKL_THREADING_LAYER=GNU \
+    HOME=/home/user \
     PATH=/home/user/.local/bin:/usr/local/cuda/bin:$PATH \
-    PYTHONPATH=$HOME/app \
+    PYTHONPATH=/home/user/app \
     PYTHONUNBUFFERED=1 \
     GRADIO_ALLOW_FLAGGING=never \
     GRADIO_NUM_PORTS=1 \
@@ -43,32 +17,59 @@ ENV HOME=/home/user \
     CUDA_TOOLKIT_ROOT_DIR=/usr/local/cuda \
     CUDA_HOME=/usr/local/cuda \
     LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH \
-    TORCH_CUDA_ARCH_LIST="8.6"
+    TORCH_CUDA_ARCH_LIST="8.6" \
+    CMAKE_PREFIX_PATH="/usr/local/cuda:/usr/lib/x86_64-linux-gnu/"
+
+# Install Python 3.9 and necessary dependencies
+RUN apt-get update && \
+    apt-get install -y --no-install-recommends \
+    python3.9 python3.9-dev python3.9-venv \
+    git wget libgl1-mesa-glx libglib2.0-0 ffmpeg libx264-dev pkg-config \
+    build-essential cmake g++ software-properties-common gnupg2 libopenblas-dev && \
+    rm -rf /var/lib/apt/lists/*
 
-# Set CUDA & cuDNN paths
-ENV CMAKE_PREFIX_PATH="/usr/local/cuda:/usr/lib/x86_64-linux-gnu/"
+# Set Python 3.9 as the default python3 and install pip
+RUN update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.9 1 && \
+    update-alternatives --set python3 /usr/bin/python3.9 && \
+    curl -sS https://bootstrap.pypa.io/get-pip.py | python3.9 && \
+    ln -sf /usr/bin/python3.9 /usr/bin/python && \
+    ln -sf /usr/local/bin/pip /usr/bin/pip3
+
+# Add NVIDIA package repositories and install CUDA Toolkit and cuDNN
+RUN wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/cuda-keyring_1.0-1_all.deb && \
+    dpkg -i cuda-keyring_1.0-1_all.deb && \
+    apt-get update && \
+    apt-get install -y --no-install-recommends cuda-toolkit-11-7 libcudnn8 libcudnn8-dev && \
+    rm -rf /var/lib/apt/lists/*
+
+# Create a non-root user
+RUN useradd -m -u 1000 user
+
+# Switch to the new user
+USER user
 
-# Set working directory
-WORKDIR $HOME/app
+# Set the working directory
+WORKDIR /home/user/app
 
 # Clone the repository
-RUN git clone -b dev https://github.com/fffiloni/dreamtalk $HOME/app 
+RUN git clone -b dev https://github.com/fffiloni/dreamtalk .
 
 # Download model checkpoints
-RUN wget https://huggingface.co/camenduru/dreamtalk/resolve/main/damo/dreamtalk/checkpoints/denoising_network.pth -O $HOME/app/checkpoints/denoising_network.pth
-RUN wget https://huggingface.co/camenduru/dreamtalk/resolve/main/damo/dreamtalk/checkpoints/renderer.pt -O $HOME/app/checkpoints/renderer.pt
+RUN mkdir -p checkpoints && \
+    wget https://huggingface.co/camenduru/dreamtalk/resolve/main/damo/dreamtalk/checkpoints/denoising_network.pth -O checkpoints/denoising_network.pth && \
+    wget https://huggingface.co/camenduru/dreamtalk/resolve/main/damo/dreamtalk/checkpoints/renderer.pt -O checkpoints/renderer.pt
 
-# Upgrade pip & install dependencies
+# Copy and install Python dependencies
 COPY requirements_HF.txt .
-RUN pip install -r requirements_HF.txt
-RUN pip install --no-cache-dir "moviepy<2" gradio
+RUN pip install --no-cache-dir -r requirements_HF.txt && \
+    pip install --no-cache-dir "moviepy<2" gradio
 
-# Copy application file
+# Copy the application file
 COPY app.py .
 
 # Set CUDA environment variables
-ENV CUDA_DEVICE_ORDER=PCI_BUS_ID
-ENV CUDA_VISIBLE_DEVICES=0
+ENV CUDA_DEVICE_ORDER=PCI_BUS_ID \
+    CUDA_VISIBLE_DEVICES=0
 
 # Run the application
-CMD ["python", "app.py"]
+CMD ["python", "app.py"]