Dockerfile

FROM docker.io/nvidia/cuda:12.3.2-cudnn9-devel-ubuntu22.04

# Set environment variables
ENV DEBIAN_FRONTEND=noninteractive
ENV PYTHONUNBUFFERED=1
ENV PYTHONDONTWRITEBYTECODE=1
ENV TF_FORCE_GPU_ALLOW_GROWTH=true

# Install system dependencies
RUN apt-get update && apt-get install -y --no-install-recommends \
    git \
    wget \
    curl \
    ca-certificates \
    python3 \
    python3-pip \
    python3-dev \
    ffmpeg \
    libsm6 \
    libxext6 \
    libgl1-mesa-glx \
    && apt-get clean \
    && rm -rf /var/lib/apt/lists/*

# Set working directory
WORKDIR /app

# Copy requirements but modify TensorFlow version
COPY requirements.txt /app/
# Use TensorFlow 2.15.0 which has better compatibility with newer CUDA versions
RUN sed -i 's/tensorflow==2.18.0/tensorflow==2.15.0/' /app/requirements.txt

# Install Python dependencies
RUN pip3 install --no-cache-dir --upgrade pip setuptools wheel
RUN pip3 install --no-cache-dir -r requirements.txt

# Install compatible OpenCV
RUN pip3 install --no-cache-dir opencv-python-headless opencv-contrib-python-headless

# Copy application code
COPY . /app/

# Create a more robust CPU fallback implementation
RUN echo 'import tensorflow as tf\n\
import os\n\
import sys\n\
\n\
# Set TensorFlow logging level\n\
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"\n\
\n\
# Function to modify FILM.py for CPU fallback\n\
def ensure_cpu_fallback():\n\
    # Check if we have a GPU available and supported\n\
    try:\n\
        gpus = tf.config.list_physical_devices("GPU")\n\
        if len(gpus) > 0:\n\
            for gpu in gpus:\n\
                tf.config.experimental.set_memory_growth(gpu, True)\n\
            print(f"Available GPUs: {len(gpus)}")\n\
        else:\n\
            print("No GPUs found, will run on CPU only")\n\
    except Exception as e:\n\
        print(f"Error setting up GPU: {e}")\n\
\n\
# Call the function\n\
ensure_cpu_fallback()\n\
' > /app/tf_setup.py

# Modify FILM.py to use CPU if GPU fails
RUN if [ -f "/app/FILM.py" ]; then \
    # Import our setup at the top of the file\
    sed -i '1s/^/import sys\nimport os\nimport tensorflow as tf\nfrom tf_setup import *\n/' /app/FILM.py && \
    # Add try-except around model call\
    sed -i '/def __call__/a\        try:' /app/FILM.py && \
    sed -i '/result = self._model/i\            # Try with GPU' /app/FILM.py && \
    sed -i '/result = self._model/a\        except Exception as e:\n            print(f"GPU inference failed: {e}, falling back to CPU")\n            with tf.device("/cpu:0"):\n                result = self._model(inputs, training=False)' /app/FILM.py; \
fi

# Set environment variables for GPU compatibility
ENV LD_LIBRARY_PATH=/usr/lib/x86_64-linux-gnu:/usr/local/cuda/lib64
ENV PATH=/usr/local/cuda/bin:${PATH}
ENV TF_FORCE_GPU_ALLOW_GROWTH=true

RUN ldconfig -p | grep cudnn

# Expose port for Streamlit
EXPOSE 8501

# Create a startup script that ensures proper execution
RUN echo '#!/bin/bash\n\
# Ensure library paths are set correctly\n\
export LD_LIBRARY_PATH=/usr/lib/x86_64-linux-gnu:/usr/local/cuda/lib64\n\
# Run the app\n\
exec streamlit run app.py --server.port=8501 --server.address=0.0.0.0\n\
' > /app/start.sh && chmod +x /app/start.sh

# Use the startup script
CMD ["/app/start.sh"]