Update Dockerfile

Dockerfile

@@ -1,6 +1,6 @@
 FROM nvidia/cuda:12.1.0-cudnn8-devel-ubuntu22.04
-# Non-interactive APT installs
-# Avoid any apt interactive prompts
+
+# Avoid apt interactive prompts
 ENV DEBIAN_FRONTEND=noninteractive \
     PATH=/opt/conda/bin:$PATH \
     CUDA_HOME=/usr/local/cuda
@@ -26,11 +26,10 @@ RUN rm -f /opt/conda/.condarc && \
     conda config --add channels conda-forge && \
     conda config --set channel_priority strict
 
-# This prevents Conda from auto-adding defaults during env creation
 ENV CONDA_FORGE=1 \
     CONDA_AUTO_UPDATE_CONDA=false
 
-# Create environment using ONLY conda-forge
+# Create env with Python 3.10
 RUN conda create --override-channels -n appenv -c conda-forge python=3.10 ffmpeg=7 -y && conda clean -afy
 SHELL ["conda", "run", "-n", "appenv", "/bin/bash", "-c"]
 
@@ -38,10 +37,15 @@ SHELL ["conda", "run", "-n", "appenv", "/bin/bash", "-c"]
 WORKDIR /workspace
 COPY requirements.txt .
 
-# Install pip dependencies
-RUN pip install --upgrade pip && pip install --no-cache-dir -r requirements.txt
+# Pre-install torch so flash-attn doesn't fail
+RUN pip install --upgrade pip && \
+    pip install torch==2.6.0 torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu121
+
+# Install the rest of the dependencies
+RUN pip install --no-cache-dir -r requirements.txt
 
-# Copy rest of code
+# Copy project code
 COPY . .
 
+# Run app
 CMD ["conda", "run", "--no-capture-output", "-n", "appenv", "python", "finetune/gradio_app.py"]