Update nightly to install flax at head

Author: Rohan-Bierneni

maxdiffusion_jax_ai_image_tpu.Dockerfile

@@ -18,5 +18,9 @@ COPY . .
 # Install Maxdiffusion Jax AI Image requirements
 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
+RUN 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