# Use official PyTorch with CUDA 12.1 (works with flash-attn)
FROM pytorch/pytorch:2.3.0-cuda12.1-cudnn8-devel

ENV DEBIAN_FRONTEND=noninteractive
ENV OMP_NUM_THREADS=4
ENV DISABLE_TRITON=1
ENV ACCELERATE_USE_DEEPSPEED=0
ENV TRANSFORMERS_VERBOSITY=info
ENV PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
ENV FLASH_ATTENTION_FORCE=1

# Install system dependencies
RUN apt-get update && apt-get install -y \
    git wget curl build-essential python3-dev \
    ffmpeg libsm6 libxext6 libgl1-mesa-glx \
    && rm -rf /var/lib/apt/lists/*

# Upgrade pip first
RUN pip install --upgrade pip

# Copy requirements (without flash-attn)
WORKDIR /app
COPY requirements.txt /app/requirements.txt
RUN grep -v "flash-attn" requirements.txt > requirements-clean.txt

# Install all Python deps except flash-attn
RUN pip install --no-cache-dir -r requirements-clean.txt

# Install flash-attn last to ensure Torch is ready
RUN pip install --no-build-isolation flash-attn==2.8.2

# Copy application
COPY . /app

# Expose Gradio
EXPOSE 7860

# Default command to launch your app
CMD ["python", "app.py"]