You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
Copy file name to clipboardExpand all lines: docs/getting_started/run_maxdiffusion_via_xpk.md
+20-1Lines changed: 20 additions & 1 deletion
Original file line number
Diff line number
Diff line change
@@ -31,7 +31,7 @@ gcloud auth login
31
31
32
32
* Run this command to configure docker to use docker-credential-gcloud for GCR registries:
33
33
```
34
-
gcloud auth configure-docker
34
+
gcloud auth configure-docker us-docker.pkg.dev
35
35
```
36
36
37
37
* Test the installation by running
@@ -61,6 +61,25 @@ after which log out and log back in to the machine.
61
61
# Default will pick stable versions of dependencies
62
62
bash docker_build_dependency_image.sh
63
63
```
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
You can find a list of available JAX Stable Stack base images [here](https://pantheon.corp.google.com/artifacts/docker/cloud-tpu-images/us/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
+
64
83
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`.
65
84
66
85
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