pin flax install

Author: Rohan-Bierneni

maxdiffusion_jax_ai_image_tpu.Dockerfile

@@ -19,7 +19,7 @@ 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
+RUN pip install --upgrade --force-reinstall git+https://github.com/google/flax.git@ef78d6584623511746be4824965cdef42b464583
 RUN pip install "fsspec==2025.10.0"
 
 # Run the script available in JAX-AI-Image base image to generate the manifest file