FROM nvidia/cuda:12.1.0-cudnn8-devel-ubuntu22.04

# Avoid apt interactive prompts
ENV DEBIAN_FRONTEND=noninteractive \
    PATH=/opt/conda/bin:$PATH \
    CUDA_HOME=/usr/local/cuda

# Install system packages
RUN apt-get update && apt-get install -y \
    wget git build-essential \
    libgl1-mesa-glx libgl1-mesa-dri xvfb ffmpeg \
    && rm -rf /var/lib/apt/lists/*

# Install Miniconda
RUN wget --quiet https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh -O /tmp/miniconda.sh && \
    bash /tmp/miniconda.sh -b -p /opt/conda && \
    rm /tmp/miniconda.sh && \
    conda clean -afy

# Remove ALL Anaconda default channels (system + user configs)
RUN rm -f /opt/conda/.condarc && \
    conda config --system --remove-key channels || true && \
    conda config --remove-key channels || true && \
    conda config --system --add channels conda-forge && \
    conda config --system --set channel_priority strict && \
    conda config --add channels conda-forge && \
    conda config --set channel_priority strict

# Ensure build tools are present
RUN apt-get update && apt-get install -y --no-install-recommends \
      ninja-build build-essential cmake git python3-dev && \
      rm -rf /var/lib/apt/lists/*

# Install packaging helpers
RUN pip install --upgrade pip setuptools wheel packaging ninja

# Uninstall any prebuilt flash-attn
ENV CONDA_FORGE=1 \
    CONDA_AUTO_UPDATE_CONDA=false

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"]

# Copy requirements
WORKDIR /workspace
COPY requirements.txt .

# Pre-install torch so flash-attn doesn't fail
# RUN pip install --upgrade pip && \
#    pip install torch==2.6 torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu121

RUN pip install torch==2.5.1+cu121 torchvision --index-url https://download.pytorch.org/whl/cu121
RUN pip install flash-attn

# Build flash-attn from source (matches PyTorch ABI)
# RUN pip install --no-build-isolation --force-reinstall flash-attn

RUN pip install -r requirements.txt
# Copy project code
COPY . .

# Run app
CMD ["conda", "run", "--no-capture-output", "-n", "appenv", "python", "finetune/gradio_app.py"]