FROM nvidia/cuda:11.8.0-runtime-ubuntu22.04

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

WORKDIR /app

# Install PyTorch and other dependencies
RUN pip3 install --upgrade pip
RUN pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
RUN pip3 install runpod transformers peft accelerate bitsandbytes gradio

# Copy your handler
COPY handler.py .

# Expose port for Gradio (if needed)
EXPOSE 7860

# Run handler
CMD [ "python3", "-u", "handler.py" ]