Update dockerfile 2

Dockerfile

@@ -1,11 +1,10 @@
-FROM nvidia/cuda:12.2.0-cudnn8-devel-ubuntu22.04
+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
-ENV XLA_FLAGS="--xla_gpu_enable_fast_min_max"
 
 # Install system dependencies
 RUN apt-get update && apt-get install -y --no-install-recommends \
@@ -26,38 +25,72 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
 # Set working directory
 WORKDIR /app
 
-# Copy requirements.txt but modify TensorFlow version
+# 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 with CUDA support
-RUN pip3 install --no-cache-dir opencv-python-headless==4.5.5.64 opencv-contrib-python-headless==4.5.5.64
+# Install compatible OpenCV
+RUN pip3 install --no-cache-dir opencv-python-headless opencv-contrib-python-headless
 
 # Copy application code
 COPY . /app/
 
-# Ensure script files have the right permissions
-RUN chmod +x /app/run_morphing.py
-RUN chmod +x /app/FILM.py
+# 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
 
-# Add FILM.py fallback to CPU if CUDA fails
-RUN sed -i '/def __call__/a\        try:' /app/FILM.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("GPU inference failed, falling back to CPU")\n            # Force CPU execution\n            with tf.device("/cpu:0"):\n                result = self._model(inputs, training=False)' /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 CUDA
-ENV LD_LIBRARY_PATH=/usr/local/cuda/lib64:${LD_LIBRARY_PATH}
+# 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}
-
-# Force TensorFlow to use compatible cuDNN version
-ENV TF_CUDNN_VERSION=8
+ENV TF_FORCE_GPU_ALLOW_GROWTH=true
 
 # Expose port for Streamlit
 EXPOSE 8501
 
-# Command to run the application
-CMD ["streamlit", "run", "app.py", "--server.port=8501", "--server.address=0.0.0.0"]
+# 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"]