Spaces:
Running
Running
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 # Force CPU execution\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/local/cuda/lib64:/usr/lib/x86_64-linux-gnu:${LD_LIBRARY_PATH} | |
ENV PATH=/usr/local/cuda/bin:${PATH} | |
ENV TF_FORCE_GPU_ALLOW_GROWTH=true | |
# 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/local/cuda/lib64:/usr/lib/x86_64-linux-gnu:$LD_LIBRARY_PATH\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"] |