unify working directory and add more gpu dependency files

docker_build_dependency_image.sh

@@ -61,6 +61,11 @@ COMMIT_HASH=$(git rev-parse --short HEAD)
 echo "Building MaxDiffusion with MODE=${MODE} at commit hash ${COMMIT_HASH} . . ."
 
 if [[ ${DEVICE} == "gpu" ]]; then
+  if [[ ${MODE} == "pinned" ]]; then
+    export BASEIMAGE=ghcr.io/nvidia/jax:base-2024-05-07
+  else
+    export BASEIMAGE=ghcr.io/nvidia/jax:base
+  fi
   docker build --network host --build-arg MODE=${MODE} --build-arg JAX_VERSION=$JAX_VERSION --build-arg DEVICE=$DEVICE --build-arg BASEIMAGE=$BASEIMAGE -f ./maxdiffusion_gpu_dependencies.Dockerfile -t ${LOCAL_IMAGE_NAME} .
 else 
   if [[ "${MODE}" == "stable_stack" ]]; then

docs/getting_started/run_maxdiffusion_via_xpk.md

@@ -80,6 +80,15 @@ after which log out and log back in to the machine.
 
     **Important Note:** The JAX Stable Stack is currently in the experimental phase. We encourage you to try it out and provide feedback.
 
+
+    #### Run MaxDiffusion on GPU
+    Default device is TPU. To run MaxDiffusion on GPU, please explicitly specify GPU When building docker image. 
+
+    ```shell
+    # Default will pick base image. 
+    bash docker_build_dependency_image.sh DEVICE=gpu 
+    ```
+
 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

maxdiffusion_dependencies.Dockerfile

@@ -45,7 +45,7 @@ ARG JAX_VERSION
 ENV ENV_JAX_VERSION=$JAX_VERSION
 
 # Set the working directory in the container
-WORKDIR /app
+WORKDIR /deps
 
 # Copy all files from local workspace into docker container
 COPY . .

maxdiffusion_jax_stable_stack_tpu.Dockerfile

@@ -7,10 +7,10 @@ ARG COMMIT_HASH
 
 ENV COMMIT_HASH=$COMMIT_HASH
 
-RUN mkdir -p /app
+RUN mkdir -p /deps
 
 # Set the working directory in the container
-WORKDIR /app
+WORKDIR /deps
 
 # Copy all files from local workspace into docker container
 COPY . .

maxdiffusion_runner.Dockerfile

@@ -2,9 +2,9 @@ ARG BASEIMAGE=maxdiffusion_base_image
 FROM $BASEIMAGE
 
 # Set the working directory in the container
-WORKDIR /app
+WORKDIR /deps
 
 # Copy all files from local workspace into docker container
 COPY . .
 
-WORKDIR /app
+WORKDIR /deps