Extension for Stable Diffusion on AWS で大規模に画像生成 AI を活用する

画像生成 AI は、テキストからリアルな画像を生成できる革新的な AI モデルです。Diffusion モデルによる技術的なブレイクスルーにより注目を集め、派生技術や関連技術の研究開発が活発化し、日進月歩で新たな技術が生まれています。画像生成 AI を利用するには、現状いくつかの方法があります。Adobe や Canva のようなアプリケーションに組み込まれている生成 AI 機能で利用する方法もあれば、Amazon Bedrocks の Amazon Titan Image Generator や、Stable Diffusion のように、API を経由してモデルをシステムに組み込んで利用する方法もあります。さらに OSS モデルや OSS アプリケーションをローカルやクラウドにデプロイして利用することもできます。

OSS アプリケーション/モデルには、利用できる派生モデルの種類が豊富で、カスタマイズ性が高く、OSS コミュニティが開発した新機能を素早く活用できるというメリットがあります。しかし、OSS モデルの推論やファインチューニングには、GPU などの計算リソースが必要となります。ローカルやオンプレミスのハードウェアリソースには限界があるため、小規模な PoC を超えて制作スタジオ等で複数人で大規模に利用する場合はクラウドのスケーラブルなコンピューティングリソースを活用することで、OSS モデルのワークロードをより効率的に実行できます。

AWS では、Extension for Stable Diffusion on AWS というソリューションを提供しています。このソリューションは OSS の Stable Diffusion Web UI の拡張機能で、モデルのファインチューニングと推論をクラウド上のリソースにオフロードすることで、複数ユーザーによる大規模な利用を可能にします。この拡張機能を使用すると、お客様は UI 上から Amazon SageMaker による Stable Diffusion のモデルトレーニング、ファインチューニング、推論を簡単に実行できます。SageMaker のマネージドな分散学習と高性能の GPU インスタンスを活用することで、必要な時だけ高パフォーマンスのリソースを利用できます。また、 CloudFormation テンプレートを使ってリソースをプロビジョニングすることで、手動設定に伴う複雑さやミスを回避できます。

本記事では、Extension for Stable Diffusion on AWS の概要、アーキテクチャ、主要機能、導入手順について詳しく解説します。

使い方

ソリューションのデプロイ

本ソリューションは「バックエンドの API を提供するミドルウェア」部分と「フロントエンドとなる Extension for Stable Diffusion on AWS 拡張機能をインストールした Stable Diffusion Web UI 」の二つのコンポーネントに分かれています。それぞれ CloudFormation でデプロイし拡張機能の画面からバックエンド API の URL を登録することで設定することができます。詳細なデプロイ手順はドキュメントおよびレポジトリをご確認ください。

ユーザーの管理

Stable Diffusion Web UI はシングルユーザー向けのアプリケーションですが、拡張機能を利用することで複数のユーザーを管理することが可能です。ロールでの権限管理が可能で、学習・チェックポイントのアップロード・推論・エンドポイントの立ち上げ・ユーザーの変更・ロールの変更などのアクションごとに権限を設定することが可能です。

モデルとエンドポイントの管理

画像生成モデルは多くのモデルが出ており、またモデルをカスタマイズできる LoRA などのモデルも多数存在します。拡張機能を利用することで、Web UI の組み込みのモデル、ローカルのチェックポイント、URL から複数のモデルをアップロードし利用することが可能です。アップロードされたモデルは S3 に配置され、必要に応じてバックエンドの SageMaker に読み込まれいつでも呼び出しが可能です。

バックエンドの SageMaker のエンドポイントも拡張機能で管理することが可能です。リアルタイムエンドポイントと非同期エンドポイントに対応しており、オートスケーリングを設定することができます。非同期エンドポイントの場合は、使わない時は自動で 0 にスケールダウンしコストを抑えて利用することが可能です。

モデルの推論

モデルの呼び出しを行ってみましょう。txt2img ページの下部に拡張機能で追加された Amazon SageMaker Inference セクションがあります。モデルを指定してリクエストを送ることでバックエンドで指定したモデルで推論が実行され生成された結果が表示されます。

また、生成されたモデルは S3 に保存されており後から見返すことが可能です。

ControlNet も利用することが可能です。以下の例では国土交通省が公開している都市 3D データ PLATEAU （CC-BY-4.0） を使用し街並みを生成しています。

Stable Diffusion Turbo や Stable Diffusion LCM などの高速化技術も利用可能です。

モデルの学習

また、データセットを登録し LoRA でファインチューニングすることが可能です。まずは、Dataset Management ページでローカルから画像をアップロードしデータセットを作成します。以下の例では AWS のプロダクションチームが作成した CG アニメーションの背景を学習させています。

Train Management ページでモデルと作成されたデータセットを指定しモデルのファインチューニングジョブを作成することができます。学習は裏側では SageMaker Training ジョブで OSS ライブラリの kohya-ss/sd-scripts を使用しています。パラメータなども指定することが可能です。

学習結果および途中のチェックポイントは S3 に自動でアップロードされ、推論で利用することができます。Amazon SageMaker Inference セクションの LoRA のドロップダウンで選択すると反映されます。以下の例では、学習させた CG アニメーション風のスタイルで生成されていることが確認できます。

まとめ

この記事では AWS が開発している Stable Diffusion Web UI の拡張機能 Extension for Stable Diffusion on AWS について紹介しました。このソリューションにより、複数ユーザーによる大規模な利用が可能になり、必要な時だけ高パフォーマンスのリソースを活用してモデルの学習、推論で利用することが可能です。ロールベースのアクセス権限管理、複数モデルの管理、リアルタイムおよび非同期の推論エンドポイント作成、ControlNet や LCM 等の高度な生成オプション、データセットの登録とLoRAによる転移学習などクリエイター向けに必要な機能と複数ユーザーで利用するための機能を備えており、これからも継続的な開発により新たな機能が追加されていく予定です。

ぜひ実際に Extension for Stable Diffusion on AWS をデプロイし、様々な機能を試してみてください。まずは小規模な環境で評価し、本番環境での運用を検討するとよいでしょう。このブログで紹介した Extension for Stable Diffusion on AWS のレポジトリとドキュメントなども参考にしてください。

著者プロフィール

前川 泰毅 (Taiki Maekawa) は AWS Japan のソリューションアーキテクトでメディア領域のお客様中心にアーキテクチャ設計や構築を支援しています。機械学習領域を得意としておりソリューションやサンプルコードの作成を行っています。