solving for those nightlys that no longer exist lol

Dockerfile

@@ -1,23 +1,20 @@
-# thecollabagepatch/magenta:latest
+# thecollabagepatch/magenta:latest - Duplicable x86_64 build
 FROM nvidia/cuda:12.6.2-cudnn-runtime-ubuntu22.04
 
-# CUDA libs present + on loader path
+# Ensure CUDA libraries are on loader path
 RUN apt-get update && apt-get install -y --no-install-recommends \
     cuda-libraries-12-4 && rm -rf /var/lib/apt/lists/*
 ENV LD_LIBRARY_PATH=/usr/local/cuda/lib64:/usr/local/cuda-12.4/lib64:/usr/local/cuda-12.4/compat:/usr/local/cuda/targets/x86_64-linux/lib:${LD_LIBRARY_PATH}
 RUN ln -sf /usr/local/cuda/targets/x86_64-linux/lib /usr/local/cuda/lib64 || true
 
-# Ensure the NVIDIA repo key is present (non-interactive) and install cuDNN 9.8
+# Install cuDNN 9.8 for better compatibility
 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.*' \
@@ -26,135 +23,123 @@ RUN set -eux; \
   ldconfig; \
   rm -rf /var/lib/apt/lists/*
 
-# (optional) preload workaround if still needed
+# Performance optimizations for L40S/Ada
 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)
 ENV TF_GPU_ALLOCATOR=cuda_malloc_async
-
-# Let cuBLAS use TF32 fast path on Ada (L40S) for big GEMMs
 ENV TF_ENABLE_CUBLAS_TF32=1 NVIDIA_TF32_OVERRIDE=1
 
 ENV DEBIAN_FRONTEND=noninteractive \
     PYTHONUNBUFFERED=1 \
     PIP_NO_CACHE_DIR=1 \
     TF_FORCE_GPU_ALLOW_GROWTH=true \
-    XLA_PYTHON_CLIENT_PREALLOCATE=false
+    XLA_PYTHON_CLIENT_PREALLOCATE=false \
+    JAX_PLATFORMS=""
 
-ENV JAX_PLATFORMS=""
+SHELL ["/bin/bash", "-c"]
 
-# --- OS deps ---
+# Install system dependencies
 RUN apt-get update && apt-get install -y --no-install-recommends \
-    software-properties-common curl ca-certificates git \
+    software-properties-common curl ca-certificates \
+    build-essential pkg-config git \
     libsndfile1 ffmpeg \
-    build-essential pkg-config \
     && add-apt-repository ppa:deadsnakes/ppa -y \
     && apt-get update && apt-get install -y --no-install-recommends \
-    python3.11 python3.11-venv python3.11-distutils python3-pip \
+    python3.11 python3.11-dev python3.11-venv python3-pip \
     && rm -rf /var/lib/apt/lists/*
 
-# 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!) ---
-# 1) JAX CUDA pins
-RUN python -m pip install "jax[cuda12]==0.7.1" "jaxlib==0.7.1"
-
-# 2) Lock seqio early to avoid backtracking madness
-RUN python -m pip install "seqio==0.0.11"
-
-# 3) Install Magenta RT *without* deps so we control pins
-RUN python -m pip install --no-deps 'git+https://github.com/magenta/magenta-realtime#egg=magenta_rt[gpu]'
-
-# 4) TF nightlies (MATCH DATES!)
-RUN python -m pip install \
-    "tf_nightly==2.20.0.dev20250619" \
-    "tensorflow-text-nightly==2.20.0.dev20250316" \
-    "tf-hub-nightly"
-
-# 5) tf2jax pinned alongside tf_nightly so pip doesn’t drag stable TF
-RUN python -m pip install tf2jax "tf_nightly==2.20.0.dev20250619"
-
-# 6) The rest of MRT deps + API runtime deps
-RUN python -m pip install \
-    gin-config librosa resampy soundfile \
-    google-auth google-auth-oauthlib google-auth-httplib2 \
-    google-api-core googleapis-common-protos google-resumable-media \
-    google-cloud-storage requests tqdm typing-extensions numpy==2.1.3 \
-    fastapi uvicorn[standard] python-multipart pyloudnorm
-
-# 7) Exact commits for T5X/Flaxformer as in pyproject
-RUN python -m pip install \
-    "t5x @ git+https://github.com/google-research/t5x.git@92c5b46" \
-    "flaxformer @ git+https://github.com/google/flaxformer@399ea3a"
-
-# ---- FINAL: enforce TF nightlies and clean any stable TF ----
-RUN python - <<'PY'
-import sys, sysconfig, glob, os, shutil
-# Find a writable site dir (site-packages OR dist-packages)
-cands = [sysconfig.get_paths().get('purelib'), sysconfig.get_paths().get('platlib')]
-cands += [p for p in sys.path if p and p.endswith(('site-packages','dist-packages'))]
-site = next(p for p in cands if p and os.path.isdir(p))
-
-patterns = [
-  "tensorflow", "tensorflow-*.dist-info", "tensorflow-*.egg-info",
-  "tf-nightly-*.dist-info", "tf_nightly-*.dist-info",
-  "tensorflow_text", "tensorflow_text-*.dist-info",
-  "tf-hub-nightly-*.dist-info", "tf_hub_nightly-*.dist-info",
-  "tf_keras-nightly-*.dist-info", "tf_keras_nightly-*.dist-info",
-  "tensorboard*", "tb-nightly-*.dist-info",
-  "keras*",  # remove stray keras
-  "tensorflow_hub*", "tensorflow_io*",
-]
-for pat in patterns:
-  for path in glob.glob(os.path.join(site, pat)):
-    if os.path.isdir(path): shutil.rmtree(path, ignore_errors=True)
-    else:
-      try: os.remove(path)
-      except FileNotFoundError: pass
-
-print("TF/Hub/Text cleared in:", site)
-PY
-
-# Reinstall pinned nightlies in ONE transaction
-RUN python -m pip install --no-cache-dir --force-reinstall \
-    "tf-nightly==2.20.0.dev20250619" \
-    "tensorflow-text-nightly==2.20.0.dev20250316" \
-    "tf-hub-nightly"
-
-RUN python -m pip install huggingface_hub
-
-RUN python -m pip install --no-cache-dir --force-reinstall "protobuf==4.25.3"
-
-RUN python -m pip install gradio
-
-RUN python -m pip install soxr
-
-# Switch to Spaces’ preferred user
-# Switch to Spaces’ preferred user
+# Make python3.11 default and install uv for faster package installs
+RUN ln -sf /usr/bin/python3.11 /usr/bin/python && \
+    ln -sf /usr/bin/python3.11 /usr/bin/python3 && \
+    python -m pip install --upgrade pip && \
+    python -m pip install uv
+
+# CRITICAL: Install TensorFlow FIRST to block tensorflow-cpu
+# Using generic tf-nightly (no specific date) for duplicability
+RUN uv pip install --system tf-nightly
+
+# Install JAX with CUDA support (pinned for stability)
+RUN uv pip install --system "jax[cuda12]==0.4.35" "jaxlib==0.4.35"
+
+# Install base dependencies
+RUN uv pip install --system \
+    absl-py chex gin-config numpy requests tqdm typing-extensions \
+    google-cloud-storage librosa resampy soundfile sentencepiece
+
+# Clone and install t5x WITHOUT dependencies (avoid tensorflow-cpu)
+RUN git clone https://github.com/google-research/t5x.git /t5x && \
+    cd /t5x && \
+    git checkout 92c5b467a5964d06c351c7eae4aa4bcd341c7ded && \
+    uv pip install --system --no-deps -e .
+
+# Install flaxformer without deps
+RUN git clone https://github.com/google/flaxformer.git /flaxformer && \
+    cd /flaxformer && \
+    git checkout 399ea3a && \
+    uv pip install --system --no-deps -e .
+
+# Install seqio without deps and PATCH OUT tensorflow_text
+RUN git clone https://github.com/google/seqio.git /seqio && \
+    cd /seqio && \
+    uv pip install --system --no-deps -e . && \
+    # CRITICAL FIX: Remove unused tensorflow_text import (not needed by Magenta RT)
+    sed -i '/import tensorflow_text as tf_text/d' /seqio/seqio/vocabularies.py
+
+# Install airio (t5x dependency) without deps
+RUN git clone https://github.com/google/airio.git /airio && \
+    cd /airio && \
+    uv pip install --system --no-deps -e .
+
+# Install clu without deps
+RUN git clone https://github.com/google/CommonLoopUtils.git /clu && \
+    cd /clu && \
+    uv pip install --system --no-deps -e .
+
+# Now install all remaining dependencies these packages need
+RUN uv pip install --system \
+    flax optax orbax-checkpoint \
+    fiddle cached_property tf2jax \
+    aqtp etils jestimator \
+    tensorflow-datasets tfds-nightly \
+    apache-beam pyyaml rouge-score sacrebleu scipy \
+    grain-nightly editdistance pyglove
+
+# Patch jestimator for newer JAX - PartitionSpec moved to jax.sharding
+RUN sed -i 's|from jax.experimental.pjit import PartitionSpec|from jax.sharding import PartitionSpec|g' \
+    /usr/local/lib/python3.11/dist-packages/jestimator/amos_helper.py || \
+    sed -i 's|from jax.experimental.pjit import PartitionSpec|from jax.sharding import PartitionSpec|g' \
+    /usr/lib/python3.11/dist-packages/jestimator/amos_helper.py || true
+
+# Install magenta-realtime without deps
+RUN git clone https://github.com/magenta/magenta-realtime.git /magenta-realtime-src && \
+    cd /magenta-realtime-src && \
+    uv pip install --system --no-deps -e .
+
+# API and audio processing dependencies
+RUN uv pip install --system \
+    fastapi uvicorn[standard] python-multipart \
+    pyloudnorm gradio soxr huggingface_hub
+
+# Ensure compatible protobuf version
+RUN uv pip install --system --force-reinstall "protobuf>=5.27.0"
+
+# Create HuggingFace Space user
 RUN useradd -m -u 1000 appuser
 WORKDIR /home/appuser/app
 
-# Copy from *build context* into image, owned by appuser
-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
+# Set cache directory
+ENV MAGENTA_RT_CACHE_DIR=/home/appuser/.cache/magenta_rt
+RUN mkdir -p $MAGENTA_RT_CACHE_DIR && chown -R appuser:appuser /home/appuser/.cache
+
+# Copy application files with proper ownership
+COPY --chown=appuser:appuser app.py /home/appuser/app/
+COPY --chown=appuser:appuser utils.py /home/appuser/app/
+COPY --chown=appuser:appuser jam_worker.py /home/appuser/app/
+COPY --chown=appuser:appuser one_shot_generation.py /home/appuser/app/
+COPY --chown=appuser:appuser model_management.py /home/appuser/app/
+COPY --chown=appuser:appuser documentation.html /home/appuser/app/
+COPY --chown=appuser:appuser lil_demo_540p.mp4 /home/appuser/app/
+COPY --chown=appuser:appuser magentaRT_rt_tester.html /home/appuser/app/
+COPY --chown=appuser:appuser magenta_prompts.js /home/appuser/app/
 COPY --chown=appuser:appuser docs/ /home/appuser/app/docs/
 
 USER appuser