FROM nvidia/cuda:12.3.2-cudnn9-devel-ubuntu22.04

# Set environment variables
ENV PYTHONUNBUFFERED=1 \
    DEBIAN_FRONTEND=noninteractive \
    CUDA_HOME=/usr/local/cuda \
    PATH=/usr/local/cuda/bin:$PATH \
    LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH \
    NVIDIA_VISIBLE_DEVICES=all \
    NVIDIA_DRIVER_CAPABILITIES=compute,utility \
    HF_HOME=/app/models \
    TRITON_CACHE_DIR=/tmp/triton_cache \
    XDG_CACHE_HOME=/tmp \
    NUMBA_CACHE_DIR=/tmp/numba_cache \
    TOKENIZERS_PARALLELISM=false \
    OMP_NUM_THREADS=4 \
    MKL_NUM_THREADS=4 \
    OPENBLAS_NUM_THREADS=4 \
    NUMEXPR_NUM_THREADS=4 \
    RAYON_NUM_THREADS=4 \
    HF_HUB_OFFLINE=1 \
    TRANSFORMERS_OFFLINE=1 \
    TORCH_COMPILE_DISABLE=1 \
    TRITON_DISABLE_LINE_INFO=1 \
    CUDA_LAUNCH_BLOCKING=1

# Install system dependencies
RUN apt-get update && apt-get install -y --no-install-recommends \
    python3 \
    python3-pip \
    python3-dev \
    build-essential \
    git \
    git-lfs \
    ffmpeg \
    libsndfile1 \
    curl \
    && rm -rf /var/lib/apt/lists/*

# Upgrade pip and install build tools
RUN python3 -m pip install --upgrade pip setuptools wheel uv

WORKDIR /app

# Create Numba cache directory
RUN mkdir -p /tmp/numba_cache /tmp/triton_cache && \
    chown nobody:nogroup /tmp/numba_cache /tmp/triton_cache && \
    chmod 700 /tmp/numba_cache /tmp/triton_cache

COPY requirements.txt .

# Install other requirements
RUN python3 -m uv pip install --no-cache-dir -r requirements.txt --prerelease=allow

COPY . .

# Check for LFS pointers and pull actual files if needed
RUN \
    has_lfs_pointer=false; \
    for file in $(find models -type f \( -name "*.pt" -o -name "*.bin" -o -name "*.safetensors" \) 2>/dev/null | head -10); do \
        if [ -f "$file" ] && head -1 "$file" 2>/dev/null | grep -q "version https://git-lfs.github.com/spec/v1"; then \
            has_lfs_pointer=true; \
            break; \
        fi; \
    done; \
    if [ "$has_lfs_pointer" = "true" ]; then \
        echo "Detected LFS pointers, initializing git and pulling LFS files..."; \
        git lfs install || true; \
        if [ ! -d .git ]; then \
            git init && \
            git config user.email "docker@example.com" && \
            git config user.name "Docker" && \
            git remote add origin https://huggingface.co/momfun/pumpnotfun-01-47 2>/dev/null || true; \
        fi; \
        git add -A 2>/dev/null || true; \
        git commit -m "Initial" 2>/dev/null || true; \
        git fetch origin 2>/dev/null || true; \
        git lfs fetch origin 2>/dev/null || true; \
        git lfs checkout || true; \
        echo "LFS files pulled successfully"; \
    else \
        echo "No LFS pointers detected, skipping git-lfs pull"; \
    fi

EXPOSE 8000

CMD ["python3", "server.py"]