Update Dockerfile

Dockerfile

@@ -58,19 +58,22 @@ RUN ln -sf /usr/bin/python3.11 /usr/bin/python && python -m pip install --upgrad
 # 1) Install JAX first
 RUN python -m pip install "jax[cuda12]==0.7.1" "jaxlib==0.7.1"
 
-# 2) Install seqio dependencies manually (excluding tfds-nightly to avoid protobuf conflict)
+# 2) Install seqio dependencies manually (excluding tfds-nightly and tensorflow-text)
 RUN python -m pip install \
     absl-py clu sentencepiece
 
 # 3) Install seqio without dependencies
 RUN python -m pip install --no-deps "seqio==0.0.11"
 
-# 4) Install correct protobuf version for tf-nightly
-RUN python -m pip install "protobuf>=6.31.1,<8.0.0"
+# 4) Comment out tensorflow-text import in seqio since it's not used by Magenta RT
+RUN sed -i '/import tensorflow_text as tf_text/d' /usr/local/lib/python3.11/dist-packages/seqio/vocabularies.py
 
-# 5) Install TensorFlow nightly without dependencies (to avoid protobuf downgrade)
+# 5) Install correct protobuf version for stable TensorFlow
+RUN python -m pip install "protobuf>=3.20.3,<6.0.0"
+
+# 6) Install stable TensorFlow (not nightly)
 RUN python -m pip install --no-deps \
-    "tf-nightly" \
+    "tensorflow[and-cuda]>=2.18.0" \
     "tensorflow-hub==0.16.1"
 
 # 6) Install remaining packages (consolidated, no duplicates)
@@ -92,13 +95,12 @@ RUN python -m pip install --no-deps \
 RUN python -m pip install --no-deps \
     'git+https://github.com/magenta/magenta-realtime#egg=magenta_rt[gpu]'
 
-    # 10) Remove conflicting TensorFlow, then upgrade to latest
-RUN python -m pip uninstall -y tensorflow tensorflow-io-gcs-filesystem && \
-    python -m pip install --upgrade --force-reinstall tf-nightly
-
+# 10) Remove conflicting stable tensorflow packages
+RUN python -m pip uninstall -y tensorflow tensorflow-cpu tensorflow-text || true
 
-# Also upgrade JAX to latest to ensure XLA compatibility
-RUN python -m pip install --upgrade "jax[cuda12]" jaxlib
+# Install stable TensorFlow instead of nightly (following KEHANG's recommendation)
+RUN python -m pip install --upgrade \
+    "tensorflow[and-cuda]>=2.18.0"
 
 # 11) Verify TensorFlow installation
 RUN python -c "import tensorflow as tf; print('='*50); print('TensorFlow version:', tf.__version__); print('='*50)"