# Use official PyTorch image with CUDA support
FROM pytorch/pytorch:2.1.0-cuda118-runtime

# Set working directory
WORKDIR /workspace

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

# Optional: Install CUDA Toolkit for compiling custom ops
RUN apt-get update && apt-get install -y --no-install-recommends \
    cuda-toolkit-11-8 && \
    ln -s /usr/local/cuda-11.8 /usr/local/cuda && \
    rm -rf /var/lib/apt/lists/*

ENV CUDA_HOME=/usr/local/cuda
ENV PATH=${PATH}:/usr/local/cuda/bin

# Copy requirements.txt
COPY requirements.txt .

# Install Python dependencies
RUN pip install --no-cache-dir -U pip && \
    pip install --no-cache-dir -r requirements.txt

# Copy app code
COPY . .

# Expose port
EXPOSE 7860

# Run the app
CMD ["python", "app.py"]