Update Dockerfile

Dockerfile

@@ -1,5 +1,4 @@
-
-# thecollabagepatch/magenta:latest somewhat working
+# thecollabagepatch/magenta:latest - optimized
 FROM nvidia/cuda:12.6.2-cudnn-runtime-ubuntu22.04
 
 # CUDA libs present + on loader path
@@ -12,13 +11,10 @@ RUN ln -sf /usr/local/cuda/targets/x86_64-linux/lib /usr/local/cuda/lib64 || tru
 RUN set -eux; \
   apt-get update && apt-get install -y --no-install-recommends gnupg ca-certificates curl; \
   install -d -m 0755 /usr/share/keyrings; \
-  # Refresh the *same* keyring the base source uses (no second source file)
   curl -fsSL https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/3bf863cc.pub \
     | gpg --batch --yes --dearmor -o /usr/share/keyrings/cuda-archive-keyring.gpg; \
   apt-get update; \
-  # If libcudnn is "held", unhold it so we can move to 9.8
   apt-mark unhold libcudnn9-cuda-12 || true; \
-  # Install cuDNN 9.8 for CUDA 12 (correct dev package name!)
   apt-get install -y --no-install-recommends \
       'libcudnn9-cuda-12=9.8.*' \
       'libcudnn9-dev-cuda-12=9.8.*' \
@@ -27,7 +23,7 @@ RUN set -eux; \
   ldconfig; \
   rm -rf /var/lib/apt/lists/*
 
-# (optional) preload workaround if still needed
+# Optional preload workaround if still needed
 ENV LD_PRELOAD=/usr/local/cuda/lib64/libcusparse.so.12:/usr/local/cuda/lib64/libcublas.so.12:/usr/local/cuda/lib64/libcublasLt.so.12:/usr/local/cuda/lib64/libcufft.so.11:/usr/local/cuda/lib64/libcusolver.so.11
 
 # Better allocator (less fragmentation than BFC during XLA autotune)
@@ -44,7 +40,7 @@ ENV DEBIAN_FRONTEND=noninteractive \
 
 ENV JAX_PLATFORMS=""
 
-# --- OS deps ---
+# --- OS dependencies ---
 RUN apt-get update && apt-get install -y --no-install-recommends \
     software-properties-common curl ca-certificates git \
     libsndfile1 ffmpeg \
@@ -57,7 +53,8 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
 # Make python3 => 3.11 for convenience
 RUN ln -sf /usr/bin/python3.11 /usr/bin/python && python -m pip install --upgrade pip
 
-# --- Python deps (pin order matters!) ---
+# --- Python dependencies (install order matters!) ---
+
 # 1) Install JAX first
 RUN python -m pip install "jax[cuda12]==0.7.1" "jaxlib==0.7.1"
 
@@ -95,60 +92,30 @@ 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) Verify TensorFlow installation
-RUN python -c "import tensorflow as tf; print('='*50); print('TensorFlow version:', tf.__version__); print('='*50)"
-
-RUN python -m pip install cached-property
-
-RUN python -c "import tensorflow as tf; print('='*50); print('FINAL 3TensorFlow version:', tf.__version__); print('='*50)"
-
-RUN python -m pip install 'git+https://github.com/google/airio.git@main'
-
-RUN python -c "import tensorflow as tf; print('='*50); print('FINAL 4TensorFlow version:', tf.__version__); print('='*50)"
-# 6) t5x/flaxformer with --no-deps
-RUN python -m pip install --no-deps \
-    "t5x @ git+https://github.com/google-research/t5x.git@92c5b46" \
-    "flaxformer @ git+https://github.com/google/flaxformer@399ea3a"
-RUN python -c "import tensorflow as tf; print('='*50); print('FINAL 5TensorFlow version:', tf.__version__); print('='*50)"
-# 7) magenta-rt last with --no-deps
-RUN python -m pip install --no-deps \
-    'git+https://github.com/magenta/magenta-realtime#egg=magenta_rt[gpu]'
+# 10) Upgrade to latest versions to support XlaCallModule v10
+RUN python -m pip install --upgrade --force-reinstall \
+    tf-nightly
 
-# 8) Verify final TensorFlow version
-RUN python -c "import tensorflow as tf; print('='*50); print('FINAL 6TensorFlow version:', tf.__version__); print('='*50)"
+# Also upgrade JAX to latest to ensure XLA compatibility
+RUN python -m pip install --upgrade "jax[cuda12]" jaxlib
 
-RUN python -m pip install huggingface_hub
-
-
-RUN python -m pip install gradio
-
-RUN python -m pip install soxr
+# 11) Verify TensorFlow installation
+RUN python -c "import tensorflow as tf; print('='*50); print('TensorFlow version:', tf.__version__); print('='*50)"
 
-# Switch to Spaces’ preferred user
-# Switch to Spaces’ preferred user
+# Switch to non-root user for security
 RUN useradd -m -u 1000 appuser
 WORKDIR /home/appuser/app
 
-# Copy from *build context* into image, owned by appuser
+# Copy application files
 COPY --chown=appuser:appuser app.py /home/appuser/app/app.py
-
-# NEW: shared utils + worker
 COPY --chown=appuser:appuser utils.py /home/appuser/app/utils.py
 COPY --chown=appuser:appuser jam_worker.py /home/appuser/app/jam_worker.py
-
 COPY --chown=appuser:appuser one_shot_generation.py /home/appuser/app/one_shot_generation.py
-
 COPY --chown=appuser:appuser model_management.py /home/appuser/app/model_management.py
-
 COPY --chown=appuser:appuser documentation.html /home/appuser/app/documentation.html
-
 COPY --chown=appuser:appuser lil_demo_540p.mp4 /home/appuser/app/lil_demo_540p.mp4
-
 COPY --chown=appuser:appuser magentaRT_rt_tester.html /home/appuser/app/magentaRT_rt_tester.html
-
 COPY --chown=appuser:appuser magenta_prompts.js /home/appuser/app/magenta_prompts.js
-
-# Create docs directory and copy documentation files
 COPY --chown=appuser:appuser docs/ /home/appuser/app/docs/
 
 USER appuser