Skip to content

Commit 15788d5

Browse files
committed
Adding section in the README to build Maxdiffusion with MODE=stable_stack
1 parent 33bb598 commit 15788d5

File tree

1 file changed

+20
-1
lines changed

1 file changed

+20
-1
lines changed

docs/getting_started/run_maxdiffusion_via_xpk.md

Lines changed: 20 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,7 @@ gcloud auth login
3131

3232
* Run this command to configure docker to use docker-credential-gcloud for GCR registries:
3333
```
34-
gcloud auth configure-docker
34+
gcloud auth configure-docker us-docker.pkg.dev
3535
```
3636

3737
* Test the installation by running
@@ -61,6 +61,25 @@ after which log out and log back in to the machine.
6161
# Default will pick stable versions of dependencies
6262
bash docker_build_dependency_image.sh
6363
```
64+
65+
#### New: Build MaxDiffusion Docker Image with JAX Stable Stack
66+
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.
67+
68+
###### What is JAX Stable Stack?
69+
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.
70+
71+
###### How to Use It
72+
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:
73+
74+
```
75+
# 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
76+
bash docker_build_dependency_image.sh MODE=stable_stack BASEIMAGE={{JAX_STABLE_STACK_BASEIMAGE}}
77+
```
78+
79+
You can find a list of available JAX Stable Stack base images [here](us-docker.pkg.dev/cloud-tpu-images/jax-stable-stack/tpu).
80+
81+
**Important Note:** The JAX Stable Stack is currently in the experimental phase. We encourage you to try it out and provide feedback.
82+
6483
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`.
6584
6685
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

0 commit comments

Comments
 (0)