# STEP 1: Change this line to a CUDA 11.7.1 base image
FROM nvidia/cuda:11.7.1-devel-ubuntu22.04

# Set up the environment
ENV DEBIAN_FRONTEND=noninteractive
ENV PYTHONUNBUFFERED=1

# Install Python and pip
RUN apt-get update && apt-get install -y \
    python3.10 python3-pip \
    && rm -rf /var/lib/apt/lists/*

# Set the working directory inside the container
WORKDIR /app

# Set a writable cache directory for Hugging Face models
ENV HF_HOME=/app/cache

# Create the cache directory and give it open permissions
RUN mkdir -p /app/cache && chmod 777 /app/cache

# STEP 2: Change this line to install PyTorch for CUDA 11.7
RUN pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu117

# Copy and install the rest of the Python dependencies
COPY requirements.txt .
RUN pip3 install --no-cache-dir -r requirements.txt

# Copy your main application code
COPY main.py .

# Expose the port your API will run on
EXPOSE 8000

# The command to start the Uvicorn server when the container runs
CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "8000"]