Adding trillium support for maxdiffusion tutorial

[saltysoup]

This is to add trillium support for the Serve Stable Diffusion XL (SDXL) using TPUs on GKE with MaxDiffusion by adding latest JAX + libtpu install binaries, and kubernetes deployment example to launch a trillium serving pod

[snippet-bot]

Here is the summary of changes.

You are about to add 1 region tag.

• ai-ml/maxdiffusion-tpu/serve_sdxl_v6e.yaml:14, tag gke_ai_ml_maxdiffusion_server_deployment

---

This comment is generated by snippet-bot.
If you find problems with this result, please file an issue at:
https://github.com/googleapis/repo-automation-bots/issues.
To update this comment, add snippet-bot:force-run label or use the checkbox below:

• Refresh this comment

[NimJay]

Looks good. I trust that you've done the appropriate testing.

Approving. Tag me if you'd like me to merge this.

[NimJay]

Question: Is there a reason we now install these via the Dockerfile?

[saltysoup]

@NimJay no, I've moved it into the requirements file with latest commit and can confirm the deployment and inference working.

The latest jax nightly comes with the required libtpu for Trillium, so have removed jaxlib as well.

As this is is a maxdiffusion with TPU tutorial, I've also changed to use jax[tpu] as the GPU binaries arent needed

[NimJay]

That awesome. Thanks for doing that.

Let me know (tag me) when you'd like me to merge this.

[saltysoup]

@NimJay please merge with main. Thank you