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 \
    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, pulling LFS files..."; \
        git lfs install || true; \
        if [ -d .git ]; then \
            git config user.email "docker@example.com" 2>/dev/null || true; \
            git config user.name "Docker" 2>/dev/null || true; \
            if git remote get-url origin >/dev/null 2>&1; then \
                echo "Using existing git remote origin for LFS pull..."; \
                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 "Warning: .git exists but no origin remote found, skipping LFS pull"; \
            fi; \
        else \
            echo "Warning: No .git directory found, cannot pull LFS files without git repository"; \
        fi; \
    else \
        echo "No LFS pointers detected, skipping git-lfs pull"; \
    fi

# Set offline mode after LFS pull is complete
ENV HF_HUB_OFFLINE=1 \
    TRANSFORMERS_OFFLINE=1

EXPOSE 8000

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