FROM pytorch/pytorch:2.2.1-cuda12.1-cudnn8-runtime

# Install git, git-lfs and basic dependencies
RUN apt-get update && apt-get install -y git git-lfs curl ffmpeg libsm6 libxext6 libgl1-mesa-glx xvfb

# Install Python dependencies including OpenCV
RUN pip install --no-cache-dir runpod requests opencv-python-headless gdown

# Set up environment variables for HuggingFace and torch cache
ENV HUGGINGFACE_HUB_CACHE=/workspace/cache
ENV TORCH_HOME=/workspace/models
ENV TRANSFORMERS_CACHE=/workspace/cache
ENV XDG_CACHE_HOME=/workspace/cache
ENV PYTHONUNBUFFERED=1

# Copy handler and startup script
COPY rp_handler.py /
COPY start.sh /
RUN chmod +x /start.sh

# Create directory stubs for workspace to avoid permission issues
RUN mkdir -p /workspace/app /workspace/cache /workspace/models

# Set the workspace as working directory
WORKDIR /

# Start with our script
CMD ["/start.sh"]