install flash attention

Dockerfile

@@ -2,11 +2,17 @@
 FROM ghcr.io/astral-sh/uv:0.2.12 as uv
 
 # Stage 2: Main application image
-FROM python:3.11.9-slim-bookworm
+FROM nvidia/cuda:12.1.1-cudnn8-devel-ubuntu22.04
 
 # Copy uv from first stage
 COPY --from=uv /uv /uv
 
+# Install Python, pip, venv, and system dependencies
+RUN apt-get update && \
+    apt-get install -y python3.11 python3.11-venv python3-pip ffmpeg && \
+    apt-get clean && \
+    rm -rf /var/lib/apt/lists/*
+
 # Create virtual environment with uv
 RUN --mount=type=cache,target=/root/.cache/uv \
     /uv venv /opt/venv
@@ -15,12 +21,6 @@ RUN --mount=type=cache,target=/root/.cache/uv \
 ENV VIRTUAL_ENV=/opt/venv \
     PATH="/opt/venv/bin:$PATH"
 
-# Install system dependencies
-RUN apt-get update && \
-    apt-get install -y ffmpeg && \
-    apt-get clean && \
-    rm -rf /var/lib/apt/lists/*
-
 # Create user and set permissions (required for HF Spaces)
 RUN useradd -m -u 1000 user && \
     chown -R user /opt/venv
@@ -29,7 +29,7 @@ RUN useradd -m -u 1000 user && \
 USER user
 WORKDIR /app
 
-# Set home to user's home directory
+# Set home to user's home directory and other envs
 ENV HOME=/home/user \
     PATH=/home/user/.local/bin:$PATH \
     HF_HOME=/home/user/.cache/huggingface \