BentoMLを使ったDeepFloyd IFのデプロイ

TL;DR

Hugging Faceは、オープンソースのモデル提供フレームワークであるBentoMLを使用してDeepFloyd IFをデプロイするワークフローを詳細に説明しています。この統合により、開発者は異なるGPUリソース間で推論ランナーを独立してスケーリングすることで、DeepFloyd IFの複雑でマルチステージなアーキテクチャを管理し、本番環境のパフォーマンスを最適化できます。

DeepFloyd IFのアーキテクチャと機能

DeepFloyd IFは、ピクセル空間で直接動作するオープンソースのテキストから画像へのモデルであり、Stable Diffusionのような潜在拡散モデルとは異なります。これは、凍結されたテキストエンコーダと3つのカスケードされたピクセル拡散モジュールからなるモジュラー構造を利用しています：

Stage 1: 64x64ピクセルのベース画像を生成します。
Stage 2 and 3: 画像を最終的に1024x1024ピクセルの解像度に段階的にアップスケールします。

複雑なプロンプトの言語理解を向上させるため、DeepFloyd IFはテキストエンコーダとしてT5-XXL-1.1 Large Language Model (LLM)を統合しています。

BentoMLを使った本番デプロイ

BentoMLは、Hugging Face Hubから本番対応のAIアプリケーションへモデルを移行するための統一されたフレームワークを提供します。デプロイメントプロセスは以下の5段階のライフサイクルに従います：

モデルを定義: PyTorchやTensorFlowなどのライブラリで訓練されたモデルを利用します。
モデルを保存: BentoMLローカルモデルストアにモデルを保存し、管理とサービスへのアクセスを行います。
BentoMLサービスを作成: service.pyファイルを使用してモデルをラップし、Runnersを介してサービスロジックを定義します。Runnersはスケールでのモデル推論を処理し、入出力処理のためのAPIを公開します。
Bentoをビルド: サービスとモデルを「Bento」にパッケージ化します。これはYAML構成ファイルを通じて必要なすべてのコードと依存関係を含むデプロイ可能なアーティファクトです。
Bentoをデプロイ: BentoをDockerイメージにコンテナ化してKubernetesへのデプロイを行うか、Yataiを使用してKubernetes上で自動スケーリングを行います。

技術的実装とリソース管理

DeepFloyd IFのデプロイには相当な計算リソースが必要です。推奨ハードウェアは、少なくとも2×16GB VRAMのGPUまたは1×40GB VRAMのGPUです。本番グレードのサービングでは、単一のTesla T4は推奨されません。

GPU割り当て戦略

BentoMLは、DeepFloyd IFパイプラインの各ステージのランナーを独立してスケーリングできるようにし、開発者が各ステージの特定のニーズに基づいてリソースを割り当てることを可能にします：

High VRAM (40GB+): すべてのモデルを単一のGPUで実行できます。
Dual Tesla T4 (15GB each): Stage 1を最初のGPUに割り当て、Stages 2と3を2番目のGPUに割り当てることができます。
Mixed GPU Setup: Stage 1をプライマリGPU（例：T4）に割り当て、Stages 2と3をVRAMが小さい2つの追加GPUに分散させることができます。

デプロイワークフロー

このセットアップを実装するために、開発者はimport_models.pyを使用して必要なステージをBentoML Model Storeにダウンロードし、service.pyを使用してJSONプロンプトとネガティブプロンプトを受け取り、生成された画像を返すAPIを定義します。最終的なアーティファクトはbentoml buildを使用してパッケージ化され、bentoml serveでローカルで提供したり、bentoml containerizeを使用してクラウドにデプロイすることができます。