Getting Started

September 17, 2025 ยท View on GitHub

We recommend starting with a single host first and then moving to multihost.

Getting Started: Local Development for single host

Running on Cloud TPUs

Local development is a convenient way to run MaxDiffusion on a single host. It doesn't scale to multiple hosts.

  1. Create and SSH to a single-host TPU (v6-8).
  • You can find here here the list of zones that support the v6(Trillium) TPUs
  • We recommend using the base VM image "v2-alpha-tpuv6e", which meets the version requirements: Ubuntu Version 22.04, Python 3.12 and Tensorflow >= 2.12.0
  1. Clone MaxDiffusion in your TPU VM.
git clone https://github.com/AI-Hypercomputer/maxdiffusion.git
cd maxdiffusion
  1. Within the root directory of the MaxDiffusion git repo, install dependencies by running:
# If a Python 3.12+ virtual environment doesn't already exist, you'll need to run the install command three times.
bash setup.sh MODE=stable DEVICE=tpu
  1. Active your virtual environment:
# Replace with your virtual environment name if not using this default name
venv_name="maxdiffusion_venv"
source ~/$venv_name/bin/activate

Getting Starting: Multihost development

[GKE, recommended] Running MaxDiffusion with xpk - Quick Experimentation and Production support