# Base image with CUDA support
FROM nvidia/cuda:12.8.0-cudnn-devel-ubuntu22.04 AS model-downloader

WORKDIR /app

# Install only essential system dependencies in a single RUN to reduce layers
RUN apt-get update && apt-get install -y --no-install-recommends \
    python3 \
    python3-pip \
    python3-dev \
    git \
    ffmpeg \
    curl \
    build-essential \
    && ln -s /usr/bin/python3 /usr/bin/python \
    && apt-get clean && rm -rf /var/lib/apt/lists/*

# Install Rust in a single command and clean up
RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y --no-modify-path --profile minimal \
    && rm -rf /root/.rustup/toolchains/*/share/doc

ENV PATH="/root/.cargo/bin:${PATH}"
ENV CC=/usr/bin/gcc
ENV CXX=/usr/bin/g++

# Upgrade pip and install Python dependencies in one layer
COPY requirements.txt .
RUN pip install --no-cache-dir --upgrade pip setuptools setuptools-rust torch && \
    pip install --no-cache-dir flash-attn  --no-build-isolation  && \
    pip install --no-cache-dir -r requirements.txt

# Set up user
RUN useradd -ms /bin/bash appuser && chown -R appuser:appuser /app
USER appuser