FROM pytorch/pytorch:2.5.1-cuda11.8-cudnn9-devel

ENV DEBIAN_FRONTEND=noninteractive \

PYTHONUNBUFFERED=1 \

CUDA_HOME=/usr/local/cuda \

PATH="$CUDA_HOME/bin:$PATH"

RUN apt-get update && apt-get install -y --no-install-recommends \

git \

curl \

ffmpeg \

libsm6 \

libxext6 \

&& apt-get clean && rm -rf /var/lib/apt/lists/*

WORKDIR /workspace/

COPY requirements.txt /workspace/requirements.txt

RUN pip install --upgrade pip \

&& pip install ninja \

&& MAX_JOBS=1 pip install flash-attn --no-build-isolation \

&& pip install -r requirements.txt \

&& pip install opencv-fixer==0.2.5 \

&& python -c "from opencv_fixer import AutoFix; AutoFix()"

CMD ["/bin/bash"]