[CoreML] Core ML Stable Diffusionを使ってみる
はじめに
先月、Appleが公開したStable DiffusionをApple Silicon環境で高速に動作させる仕組み（厳密にはStableDiffusionのModelをCoreML形式のModelに変換してCoreMLで動かす)を試してみました。
基本的には以下の公式の公開している内容を試してみた形(+α)となります。
- 公式のGitHub: ml-stable-diffusion
試した環境
パフォーマンスの話やOSバージョンの話もあるので環境を記しておきます。
- PC : MacBook Pro 14インチ / 2021
- 10コアCPU + 32コアGPU搭載 M1 Max
- 64 GB memory
- OS : macOS Ventura 13.0 & 13.1
環境準備
公式のREADME.mdが必要十分な内容となっているので、基本的にはそのまま実施します。
ソースコードのClone
はじめに紹介した公式のGitHub から必要なソースコード一式をclone してきます。
$ git clone https://github.com/apple/ml-stable-diffusion.git
Cloning into 'ml-stable-diffusion'...
remote: Enumerating objects: 142, done.
remote: Counting objects: 100% (38/38), done.
remote: Compressing objects: 100% (25/25), done.
remote: Total 142 (delta 19), reused 13 (delta 13), pack-reused 104
Receiving objects: 100% (142/142), 9.08 MiB | 15.02 MiB/s, done.
Resolving deltas: 100% (48/48), done.
ml-stable-diffusion というディレクトリが作られるので、以降の作業はこの下で実施します。
$ ls
ml-stable-diffusion
$ cd ml-stable-diffusion
Python 環境のインストール
conda を利用して今回の作業用の仮想環境を作成し、仮想環境下でインストールを実施します。
もし、conda がインストールされていない環境の場合は、conda のインストールからになります。
condaのインストールをする人向け
AppleのCoreML のページでは Minicondaによるセットアップを紹介していたので、Miniconda installerでmacOS Apple M1 64-bit 向けの物を利用すれば良いかと思います。
私の環境ではM1 Mac向けのTensorflow環境を以前に作った事があったため、こちらのMiniconda3-latest-MacOSX-arm64.shを利用した環境となっています。
環境のインストールは公式の記述通り、以下の流れで行います。
conda create -n coreml_stable_diffusion python=3.8 -y
conda activate coreml_stable_diffusion
conda activate を実施した時点で zsh環境では以下のように 環境がcoreml_stable_diffusionに切り替わっているのが確認できるかと思います。
(coreml_stable_diffusion) {ユーザ名} {現在のディレクトリ} $
次にpip install -e . を実行します。
pip install -e .
正常にインストールが成功している場合は、以下のようにpip listコマンドを実行した際、coremltools, diffusers, huggingface-hub といったPackageがインストール済みである事が確認出来ます。
インストールパッケージ確認
$ pip list
Package Version Editable project location
------------------------------ ---------- ---------------------------
accelerate 0.15.0
certifi 2022.9.24
charset-normalizer 2.1.1
coremltools 6.1
diffusers 0.10.2
filelock 3.8.2
huggingface-hub 0.11.1
...
環境構築としては以上となります。
Hugging Faceの準備
他のStable Diffusion のローカル環境構築の手順に漏れず、Hugging Faceのアカウント登録とユーザアクセストークンの発行が必要になります。
この辺は、参考になる記事が大量にあるかと思いますので割愛します。
ユーザアクセストークンの準備ができた後、huggingface-cli login を実行します。
$ huggingface-cli login
_| _| _| _| _|_|_| _|_|_| _|_|_| _| _| _|_|_| _|_|_|_| _|_| _|_|_| _|_|_|_|
_| _| _| _| _| _| _| _|_| _| _| _| _| _| _| _|
_|_|_|_| _| _| _| _|_| _| _|_| _| _| _| _| _| _|_| _|_|_| _|_|_|_| _| _|_|_|
_| _| _| _| _| _| _| _| _| _| _|_| _| _| _| _| _| _| _|
_| _| _|_| _|_|_| _|_|_| _|_|_| _| _| _|_|_| _| _| _| _|_|_| _|_|_|_|
To login, `huggingface_hub` now requires a token generated from https://huggingface.co/settings/tokens .
Token: (ここでユーザアクセストークンを貼り付ける)
Add token as git credential? (Y/n) y
Token is valid.
Login successful
Model のCoreML用への変換
Stable Diffusion 用のModel を以下のコマンドで変換します。<output-mlpackages-directory> の部分は適宜、出力したいディレクトリ名で置き換えてください。
python -m python_coreml_stable_diffusion.torch2coreml --convert-unet --convert-text-encoder --convert-vae-decoder --convert-safety-checker -o <output-mlpackages-directory>
デフォルトではCompVis/stable-diffusion-v1-4が変換対象として利用されます。
--model-version 使いたいモデル名と指定する事で別のモデルを変換させる事ができます。
それ以外にもチューニングのためのパラメータがいくつか存在するため、こだわりたい方は公式の説明を確認しましょう。
実行時パフォーマンスの話
せっかくApple Silicon環境で構築する以上、パフォーマンスは強い関心事項だと思います。
公式のベンチマークでも設定によって結果が違っており、かつ、設定のうち一つはこの変換処理時に設定するパラメータになっていますので、環境に合わせてパラメータの選択が必要です。
GPUパワーを期待できる環境では--attention-implementation ORIGINAL を指定して変換を実施します。
iPad等のモバイル環境や通常のApple Silicon環境では--attention-implementation SPLIT_EINSUMを利用する方が良いようです。
その場合、デフォルトでは SPLIT_EINSUMが利用されるため明示的な指定は不要です。
実行時環境 (Python or Swift) の選択の話
変換後のModelはサンプルとして提供されているPython コード or Swift コードで実行できます。
Swift コードで実施させる為には、上記のパラメータに加えて、--bundle-resources-for-swift-cli を指定します。
私の環境では以下の設定で実行を行なっています。
python -m python_coreml_stable_diffusion.torch2coreml --convert-unet --convert-text-encoder --convert-vae-decoder --convert-safety-checker --attention-implementation ORIGINAL --bundle-resources-for-swift-cli -o model_original
Torchバージョンのワーニングの話
手元の環境で実行した際は以下のようなワーニングが表示されていました。
setup.py ではtorch のバージョンを指定されていなかったので、そんな事を言われても...という気持ちになりましたが、一応動作はしていました。
Torch version 1.13.1 has not been tested with coremltools. You may run into unexpected errors. Torch 1.12.1 is the most recent version that has been tested.
気になるのであれば、pip install -e . を実行するより前にsetup.pyでバージョンを指定しておくと良いかもしれません。
install_requires=[
"coremltools>=6.1",
"diffusers[torch]",
- "torch",
+ "torch==1.12.1",
"transformers",
"scipy",
],
実行に成功すると指定したディレクトリ以下に変換されたファイルが出力されます。
- Resources (
--bundle-resources-for-swift-cliを指定した場合のみ) - Stable_Diffusion_version_CompVis_stable-diffusion-v1-4_safety_checker.mlpackage
- Stable_Diffusion_version_CompVis_stable-diffusion-v1-4_text_encoder.mlpackage
- Stable_Diffusion_version_CompVis_stable-diffusion-v1-4_unet.mlpackage
- Stable_Diffusion_version_CompVis_stable-diffusion-v1-4_vae_decoder.mlpackage
画像生成
CoreML 用に変換したModel を利用して画像を生成します。
画像生成にはPythonを利用する方法とSwiftを利用する方法の２種類があります。
<output-mlpackages-directory>の部分は先ほど生成したModel の出力パスになります。
</path/to/output/image>は任意の画像の出力先です。
各種パラメータの指定方法が異なっていますので注意してください。
Pythonスクリプト
python -m python_coreml_stable_diffusion.pipeline --prompt "a photo of an astronaut riding a horse on mars" -i <output-mlpackages-directory> -o </path/to/output/image> --compute-unit ALL --seed 93
-
--prompt: 入力文字列
--i: 利用するモデル -
-o: 出力ディレクトリ -
--compute-unit: 利用するハードウェアリソース -
--seed: シード値
`--compute-unit` について
モデル変換時に --attention-implementation ORIGINAL を指定している場合は
--compute-unit CPU_AND_GPU の指定を利用する事で高速化が見込めます。
また余談ですが、mac OSのバージョンが13.1 未満の場合、CPU_AND_GPUは動作しませんでした。(ALLは遅いですが動作します)
元々、System Requirementsが13.1なので、あえて古いOSで試す必要はないかと思いますが...
Swiftコード
--bundle-resources-for-swift-cliを指定して変換した時のみ利用可能
swift run StableDiffusionSample "a photo of an astronaut riding a horse on mars" --resource-path <output-mlpackages-directory>/Resources/ --seed 93 --output-path </path/to/output/image>
-
--image-count: 生成枚数の指定 (Swift版にのみ存在)
生成に成功した場合、</path/to/output/image> 以下に画像ファイルが出力されます。
また、手元で試した出力時間は以下の通りです。
やっぱり ALL よりは CPU & GPU の設定の方が速いです。
| 生成方法 | 計測1回目 | 計測2回目 | 計測3回目 |
|---|---|---|---|
| Python(ALL) | 2:26.81 | 2:15.08 | 2:13.27 |
| Python(CPU_AND_GPU) | 52.809 | 52.869 | 52.749 |
| Swift(cpuAndGPU) | 49.165 | 40.328 | 39.879 |
任意のモデルを使ってみる
modelの変換の際、--model-version を指定する事でhuggingface-hubにある他のStable Diffusion モデルを利用する事もできます。
ただし、私が試してみた限りではいくつかのハマりポイントがありました。
Python Scriptを使う場合は実行時も--model-versionが必要
Swift の場合は不要ですが、Python Script を利用して画像生成する場合は、変換時に指定した--model-versionと同じパラメータを指定する必要があります。
画像生成時に生成サイズを指定することができない
これは 変換ツール が最適化のために変換時に利用したサイズ(元モデルのsample size)用の CoreML Model を生成するからの模様です。
Q7: How do I generate images with different resolutions using the same Core ML models?
A7: The current version of python_coreml_stable_diffusion does not support single-model multi-resolution out of the box. However, developers may fork this project and leverage the flexible shapes support from coremltools to extend the torch2coreml script by using coremltools.EnumeratedShapes. Note that, while the text_encoder is agnostic to the image resolution, the inputs and outputs of vae_decoder and unet models are dependent on the desired image resolution.
公式の回答にもありますが、変換スクリプトを拡張することで対応が出来るらしいですが、今回はそこまで試してはいません。
--latent-h, --latent-w がうまく動作しない
上記の改修を行わずともtorch2coreml scriptには、サイズを指定するパラメータがありましたので変換時にこちらを利用してみました。
python -m python_coreml_stable_diffusion.torch2coreml --convert-unet --convert-text-encoder --convert-vae-decoder --convert-safety-checker --bundle-resources-for-swift-cli --attention-implementation ORIGINAL --latent-h 64 --latent-w 64 -o {出力先} --model-version {利用するモデル}
指定するサイズは、画素数の1/8のようなので、512 x 512 を生成するには --latent-h 64 --latent-w 64 の指定が必要です。
ただし、試してみると上記の指定した値は vae_decoder のinput_shapeにしか利用されておらず、それ以降のモデルの変換には利用されていません。
これが意図通りなのかどうかは分かりませんが、モデルによっては画像生成時にサイズの不整合でエラーが発生してしまったので、以下のように修正を行う必要がありました。
diff ファイル
diff --git a/python_coreml_stable_diffusion/torch2coreml.py b/python_coreml_stable_diffusion/torch2coreml.py
index 6d6c2fa..a7a5696 100644
--- a/python_coreml_stable_diffusion/torch2coreml.py
+++ b/python_coreml_stable_diffusion/torch2coreml.py
@@ -484,8 +484,8 @@ def convert_unet(pipe, args):
sample_shape = (
batch_size, # B
pipe.unet.config.in_channels, # C
- pipe.unet.config.sample_size, # H
- pipe.unet.config.sample_size, # W
+ args.latent_h or pipe.unet.config.sample_size, # H
+ args.latent_w or pipe.unet.config.sample_size, # w
)
if not hasattr(pipe, "text_encoder"):
@@ -619,11 +619,16 @@ def convert_safety_checker(pipe, args):
f"`safety_checker` already exists at {out_path}, skipping conversion."
)
return
+
+ VAE_DECODER_UPSAMPLE_FACTOR = 8
+
+ sample_size_h = None if args.latent_h is None else args.latent_h * VAE_DECODER_UPSAMPLE_FACTOR
+ sample_size_w = None if args.latent_w is None else args.latent_w * VAE_DECODER_UPSAMPLE_FACTOR
sample_image = np.random.randn(
1, # B
- pipe.vae.config.sample_size, # H
- pipe.vae.config.sample_size, # w
+ sample_size_h or pipe.vae.config.sample_size, # H
+ sample_size_w or pipe.vae.config.sample_size, # w
3 # C
).astype(np.float32)
Discussion