FROM huggingface/transformers-pytorch-gpu:latest

# Set basic environment variables
ENV PYTHONUNBUFFERED=1

# Install system dependencies
RUN apt-get update && apt-get install -y --no-install-recommends \
    curl \
    git \
    && rm -rf /var/lib/apt/lists/*

# Set working directory
WORKDIR /app

# Install Python requirements
RUN pip install --no-cache-dir gradio==3.38.0 pillow numpy

# Copy diagnostic script
COPY gpu_test.py /app/

# Add a script to check GPU status at startup
RUN echo '#!/bin/bash \n\
echo "==== GPU DIAGNOSTICS STARTUP CHECKS ====" \n\
echo "Checking NVIDIA driver and CUDA:" \n\
nvidia-smi || echo "nvidia-smi failed - GPU may not be properly configured" \n\
echo "Current GPU-related environment variables:" \n\
echo "CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES}" \n\
echo "NVIDIA_VISIBLE_DEVICES=${NVIDIA_VISIBLE_DEVICES}" \n\
echo "==== STARTING APPLICATION ====" \n\
exec "$@"' > /entrypoint.sh && \
chmod +x /entrypoint.sh

# Make port 7860 available for the app
EXPOSE 7860

# Use our entrypoint script to check GPU status before starting the app
ENTRYPOINT ["/entrypoint.sh"]

# Start the application
CMD ["python", "gpu_test.py"]