Installation

MaxDiffusion can be installed on Cloud TPUs, GPUs, or local machines. We recommend starting with a single TPU host for development before scaling to multi-host configurations.

System requirements

Operating System: Ubuntu 22.04 or later
Python: 3.12 or later
TensorFlow: 2.12.0 or later
Hardware: Cloud TPU v4/v5p/v6e or NVIDIA GPU with CUDA support

Quick install

Cloud TPU
GPU
Local (CPU)

Create a TPU VM

Create and SSH into a single-host TPU VM:

gcloud compute tpus tpu-vm create maxdiffusion-vm \
  --zone=us-central2-b \
  --accelerator-type=v6e-8 \
  --version=v2-alpha-tpuv6e

gcloud compute tpus tpu-vm ssh maxdiffusion-vm --zone=us-central2-b

For TPU v6e (Trillium), use the v2-alpha-tpuv6e VM image which includes Ubuntu 22.04, Python 3.12, and TensorFlow 2.12+.

Clone the repository

git clone https://github.com/AI-Hypercomputer/maxdiffusion.git
cd maxdiffusion

Run setup script

Install MaxDiffusion and dependencies:

bash setup.sh MODE=stable DEVICE=tpu

If using Python 3.12+ for the first time, you may need to run the setup script three times to complete the installation.

Activate virtual environment

venv_name="maxdiffusion_venv"
source ~/$venv_name/bin/activate

Verify installation

Test your installation by generating an image:

python -m src.maxdiffusion.generate \
  src/maxdiffusion/configs/base21.yml \
  run_name="test" \
  prompt="A magical castle in the middle of a forest"

Clone the repository

git clone https://github.com/AI-Hypercomputer/maxdiffusion.git
cd maxdiffusion

Install JAX for GPU

pip install -U "jax[cuda12]"

Install MaxDiffusion

pip install -r requirements.txt
pip install .

Install Transformer Engine (optional)

For fused attention on GPU:

pip install --upgrade torch torchvision
pip install "transformer_engine[jax]"

Verify installation

Test with SDXL on GPU:

python -m src.maxdiffusion.generate_sdxl \
  src/maxdiffusion/configs/base_xl.yml \
  hardware=gpu \
  run_name="test"

GPU support is not actively maintained. For production workloads, we recommend Cloud TPUs.

Clone the repository

git clone https://github.com/AI-Hypercomputer/maxdiffusion.git
cd maxdiffusion

Create virtual environment

python3 -m venv maxdiffusion_venv
source maxdiffusion_venv/bin/activate

Install dependencies

pip install -r requirements.txt
pip install .

CPU installation is primarily for development and testing. Training and inference will be significantly slower than on TPU or GPU.

TPU VM images

For different TPU generations, use the appropriate VM image:

TPU Type	Recommended VM Image	Zones
v6e-8 (Trillium)	`v2-alpha-tpuv6e`	us-central2-b, us-south1-a
v5p-8	`v2-alpha-tpuv5-lite`	us-east5-a, us-west4-a
v4-8	`tpu-ubuntu2204-base`	us-central2-b, europe-west4-a

See the Cloud TPU regions and zones documentation for the full list of available zones.

Development installation

For contributing to MaxDiffusion or modifying the source code:

git clone https://github.com/AI-Hypercomputer/maxdiffusion.git
cd maxdiffusion

# Install in editable mode
pip install -e .

# Install development dependencies
pip install pytest pylint pyink

Optional dependencies

For training

pip install accelerate datasets tensorboard

For video models

pip install opencv-python
apt-get update && apt-get install ffmpeg libsm6 libxext6 -y

For testing

pip install pytest pytest-timeout pytest-xdist

Troubleshooting

ImportError: cannot import JAX

Make sure JAX is installed for your target device:

# For TPU
pip install jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html

# For GPU
pip install -U "jax[cuda12]"

Python version mismatch

MaxDiffusion requires Python 3.12+. Check your version:

python --version

If using an older version, create a new virtual environment with Python 3.12.

Out of disk space during installation

MaxDiffusion models can be large. Consider attaching an external disk:

# Create and attach a 500GB disk
gcloud compute disks create maxdiffusion-disk \
  --size=500GB \
  --zone=us-central2-b

# Attach to your TPU VM
# See: https://cloud.google.com/tpu/docs/attach-durable-block-storage

Setup script fails on first run

For Python 3.12+ virtual environments, you may need to run the setup script multiple times:

bash setup.sh MODE=stable DEVICE=tpu
bash setup.sh MODE=stable DEVICE=tpu
bash setup.sh MODE=stable DEVICE=tpu

Next steps

Quickstart

Generate your first image

Single host deployment

Set up local development

Multi-host deployment

Scale to TPU Pods

Training guide

Start training models

Getting Started

Core Concepts

Training

Inference

Advanced Features

Deployment

Guides

Installation

Installation

System requirements

Quick install

TPU VM images

Development installation

Optional dependencies

For training

For video models

For testing

Troubleshooting

Next steps

Quickstart

Single host deployment

Multi-host deployment

Training guide

Build docs developers (and LLMs) love

Getting Started

Core Concepts

Training

Inference

Advanced Features

Deployment

Guides

​Installation

​System requirements

​Quick install

​TPU VM images

​Development installation

​Optional dependencies

​For training

​For video models

​For testing

​Troubleshooting

​Next steps

Quickstart

Single host deployment

Multi-host deployment

Training guide

Build docs developers (and LLMs) love

Installation

System requirements

Quick install

TPU VM images

Development installation

Optional dependencies

For training

For video models

For testing

Troubleshooting

Next steps