このチュートリアルでは、MaxDiffusion を備えた Google Kubernetes Engine（GKE）上で Tensor Processing Unit（TPU）を使用して SDXL 画像生成モデルを提供する方法について説明します。このチュートリアルでは、Hugging Face からモデルをダウンロードし、MaxDiffusion を実行するコンテナを使用して Autopilot クラスタまたは Standard クラスタにデプロイします。
このガイドは、AI / ML ワークロードをデプロイして提供する際に、マネージド Kubernetes での詳細な制御、カスタマイズ、拡張性、復元力、ポータビリティ、費用対効果が求められる場合の出発点として適しています。ML モデルを費用対効果の高い方法で迅速に構築して提供するために、統合されたマネージド AI プラットフォームが必要な場合は、Vertex AI デプロイ ソリューションをお試しになることをおすすめします。
背景
GKE で TPU を使用して MaxDiffusion で SDXL をサービングすることで、マネージド Kubernetes のメリット（費用効率、拡張性、高可用性など）をすべて活用した、本番環境対応の堅牢なサービング ソリューションを構築できます。このセクションでは、このチュートリアルで使用されている重要なテクノロジーについて説明します。
Stable Diffusion XL（SDXL）
Stable Diffusion XL（SDXL）は、推論用に MaxDiffusion がサポートする潜在的な拡散モデル（LDM）の一種です。生成 AI の場合、LDM を使用してテキストの説明から高品質の画像を生成できます。LDM は、画像検索や画像キャプションなどのアプリケーションに役立ちます。
SDXL は、シャーディング アノテーションを使用して、単一ホストまたはマルチホストの推論をサポートします。これにより、SDXL を複数のマシンでトレーニングして実行できるため、効率性が向上します。
詳細については、Stability AI リポジトリによる生成モデルと SDXL の論文をご覧ください。
TPU
TPU は、Google が独自に開発した特定用途向け集積回路（ASIC）であり、TensorFlow、PyTorch、JAX などのフレームワークを使用して構築された機械学習モデルと AI モデルを高速化するために使用されます。
GKE で TPU を使用する前に、次の学習プログラムを完了することをおすすめします。
- Cloud TPU システム アーキテクチャで、現在の TPU バージョンの可用性について学習する。
- GKE の TPU についてを確認する。
このチュートリアルでは、SDXL モデルのサービングについて説明します。GKE は、低レイテンシでプロンプトをサービングするモデルの要件に基づいて構成された TPU トポロジを使用して、単一ホストの TPU v5e ノードにモデルをデプロイします。このガイドでは、モデルで 1x1 トポロジの TPU v5e チップを使用します。
MaxDiffusion
MaxDiffusion は、Python と Jax で記述された、XLA デバイス（TPU や GPU など）で実行されるさまざまな潜在的な拡散モデルのリファレンス実装のコレクションです。MaxDiffusion は、研究環境と本番環境の両方での Diffusion プロジェクトの開始点です。
詳細については、MaxDiffusion リポジトリをご覧ください。