Install only if nightly

Author: Rohan-Bierneni

maxdiffusion_jax_ai_image_tpu.Dockerfile

@@ -19,8 +19,10 @@ COPY . .
 RUN pip install -r /deps/requirements_with_jax_ai_image.txt
 
 # TODO: Remove the flax pin and fsspec overrides once flax stable version releases
-RUN pip install --upgrade --force-reinstall git+https://github.com/google/flax.git@ef78d6584623511746be4824965cdef42b464583
-RUN pip install "fsspec==2025.10.0"
+RUN if echo "$JAX_AI_IMAGE_BASEIMAGE" | grep -q "nightly"; then \
+        echo "Nightly build detected: Installing specific Flax commit and fsspec." && \
+        pip install --upgrade --force-reinstall git+https://github.com/google/flax.git@ef78d6584623511746be4824965cdef42b464583 && \
+        pip install "fsspec==2025.10.0";
 
 # Run the script available in JAX-AI-Image base image to generate the manifest file
 RUN bash /jax-ai-image/generate_manifest.sh PREFIX=maxdiffusion COMMIT_HASH=$COMMIT_HASH