FROM pytorch/pytorch:2.5.1-cuda12.1-cudnn8-runtime

RUN apt-get update && apt-get install -y \
    git build-essential python3-dev ninja-build && \
    rm -rf /var/lib/apt/lists/*

RUN pip install --upgrade pip

COPY requirements.txt /tmp/requirements.txt
RUN pip install -r /tmp/requirements.txt

RUN pip install git+https://github.com/Dao-AILab/flash-attention.git@v2.7.3#subdirectory=csrc

COPY . /app
WORKDIR /app

CMD ["python", "app.py"]