Dockerfile

FROM pytorch/pytorch:2.0.1-cuda11.7-cudnn8-runtime

# Set environment variables
ENV DEBIAN_FRONTEND=noninteractive
ENV PYTHONUNBUFFERED=1
ENV HF_HOME=/app/.cache/huggingface
ENV TRANSFORMERS_CACHE=/app/.cache/huggingface/transformers
ENV PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:128

# Create necessary directories with proper permissions
RUN mkdir -p /app/.cache/huggingface/transformers && \
    chmod -R 777 /app

# Install system dependencies
RUN apt-get update && apt-get install -y --no-install-recommends \
    build-essential \
    git \
    curl \
    ca-certificates \
    python3-pip \
    python3-dev \
    && rm -rf /var/lib/apt/lists/*

# Create a working directory
WORKDIR /app

# Install core requirements
COPY requirements.txt .
RUN pip3 install --no-cache-dir --upgrade pip && \
    pip3 install --no-cache-dir -r requirements.txt

# Install additional dependencies specifically for InternViT
RUN pip3 install --no-cache-dir transformers==4.37.2 timm==0.9.11 accelerate==0.30.0 safetensors==0.4.1 einops

# Copy the application
COPY simple_internvit_test.py .

# Add a script to check GPU status and run the app
RUN echo '#!/bin/bash \n\
echo "Checking NVIDIA GPU status..." \n\
if ! command -v nvidia-smi &> /dev/null; then \n\
    echo "WARNING: nvidia-smi command not found. NVIDIA driver might not be installed." \n\
else \n\
    echo "NVIDIA driver found. Running nvidia-smi:" \n\
    nvidia-smi \n\
fi \n\
echo "Environment variables:" \n\
echo "CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES}" \n\
echo "NVIDIA_VISIBLE_DEVICES=${NVIDIA_VISIBLE_DEVICES}" \n\
echo "PYTORCH_CUDA_ALLOC_CONF=${PYTORCH_CUDA_ALLOC_CONF}" \n\
echo "\nStarting app..." \n\
exec "$@"' > /entrypoint.sh && \
chmod +x /entrypoint.sh

# Expose port 7860 for Gradio
EXPOSE 7860

# Use our entrypoint script
ENTRYPOINT ["/entrypoint.sh"]

# Start the application
CMD ["python3", "simple_internvit_test.py"]