diff --git a/docs/getting_started/run_maxdiffusion_via_xpk.md b/docs/getting_started/run_maxdiffusion_via_xpk.md index 21b7ff8..50f82a7 100644 --- a/docs/getting_started/run_maxdiffusion_via_xpk.md +++ b/docs/getting_started/run_maxdiffusion_via_xpk.md @@ -31,7 +31,7 @@ gcloud auth login * Run this command to configure docker to use docker-credential-gcloud for GCR registries: ``` -gcloud auth configure-docker +gcloud auth configure-docker us-docker.pkg.dev ``` * Test the installation by running @@ -61,6 +61,25 @@ after which log out and log back in to the machine. # Default will pick stable versions of dependencies bash docker_build_dependency_image.sh ``` + + #### New: Build MaxDiffusion Docker Image with JAX Stable Stack + We're excited to announce that you can build the MaxDiffusion Docker image using the JAX Stable Stack base image. This provides a more reliable and consistent build environment. + + ###### What is JAX Stable Stack? + JAX Stable Stack provides a consistent environment for MaxDiffusion by bundling JAX with core packages like `orbax`, `flax`, and `optax`, along with Google Cloud utilities and other essential tools. These libraries are tested to ensure compatibility, providing a stable foundation for building and running MaxDiffusion and eliminating potential conflicts due to incompatible package versions. + + ###### How to Use It + To build the MaxDiffusion Docker image with JAX Stable Stack, simply set the MODE to `stable_stack` and specify the desired `BASEIMAGE` in the `docker_build_dependency_image.sh` script: + + ``` + # Example bash docker_build_dependency_image.sh MODE=stable_stack BASEIMAGE=us-docker.pkg.dev/cloud-tpu-images/jax-stable-stack/tpu:jax0.4.33-rev1 + bash docker_build_dependency_image.sh MODE=stable_stack BASEIMAGE={{JAX_STABLE_STACK_BASEIMAGE}} + ``` + + You can find a list of available JAX Stable Stack base images [here](us-docker.pkg.dev/cloud-tpu-images/jax-stable-stack/tpu). + + **Important Note:** The JAX Stable Stack is currently in the experimental phase. We encourage you to try it out and provide feedback. + 3. After building the dependency image `maxdiffusion_base_image`, xpk can handle updates to the working directory when running `xpk workload create` and using `--base-docker-image`. See details on docker images in xpk here: https://github.com/google/xpk/blob/main/README.md#how-to-add-docker-images-to-a-xpk-workload