FROM nvidia/cuda:12.4.1-devel-ubuntu22.04

ENV DEBIAN_FRONTEND=noninteractive
ENV PIP_NO_CACHE_DIR=1
ENV PATH=/opt/venv/bin:$PATH

# System dependencies
RUN apt-get update && apt-get install -y --no-install-recommends \
    python3.10 \
    python3.10-venv \
    python3-pip \
    python3-dev \
    git \
    wget \
    curl \
    build-essential \
    libstdc++-11-dev \
    && apt-get clean && rm -rf /var/lib/apt/lists/*

# Virtual environment
RUN python3.10 -m venv /opt/venv
RUN pip install --upgrade pip setuptools wheel

# PyTorch 2.5.1 (CUDA 12.4)
RUN pip install torch==2.5.1 torchvision==0.20.1 torchaudio==2.5.1 \
    --index-url https://download.pytorch.org/whl/cu124

# Flash Attention 2.8.3 (pre-compiled wheel for torch 2.5 + CUDA 12.4)
RUN pip install packaging ninja psutil
RUN pip install xformers==0.0.28.post3 --index-url https://download.pytorch.org/whl/cu124
RUN pip install https://github.com/Dao-AILab/flash-attention/releases/download/v2.8.3/flash_attn-2.8.3+cu12torch2.5cxx11abiFALSE-cp310-cp310-linux_x86_64.whl

# Unsloth
RUN pip install unsloth
RUN pip install --upgrade accelerate

# Re-pin torch (unsloth may upgrade it)
RUN pip install torch==2.5.1 torchvision==0.20.1 torchaudio==2.5.1 \
    --index-url https://download.pytorch.org/whl/cu124
RUN pip install xformers==0.0.28.post3 --index-url https://download.pytorch.org/whl/cu124

# Remove torchao (incompatible with torch 2.5.1, needs torch.int1 from 2.6+)
RUN pip uninstall -y torchao || true

# Phoenix dependencies
RUN pip install timm zuko

# Gradio + utilities
RUN pip install gradio numpy Pillow huggingface_hub matplotlib

# HF cache to persistent storage (mount bucket at /data in Space settings)
ENV HF_HOME=/data/hf_cache

# Create user for HF Spaces
RUN useradd -m -u 1000 user
RUN mkdir -p /data/hf_cache && chown -R user:user /data
USER user
WORKDIR /home/user/app

# Copy application
COPY --chown=user:user . /home/user/app

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