コンテンツへスキップ

apple/ml-stable-diffusionの動作確認

タグ:

せっかくMacBook Air M2があるので、Appleが公開しているCoreML向けのStable Diffusionを動かしてみました。作業内容自体は https://github.com/apple/ml-stable-diffusion に記載されている内容そのままです。

Appleが公開しているだけあって、CoreMLを使用しての推論とSwift対応が行われています。

環境構築

作業するには以下の様なフォルダ構成にしておくのが良いです。

./ml_stable_diffusion
    ./source(GitHubからcloneしたコードを置く場所)
    ./model(CoreML向けに変換したモデルを配置する場所)
    ./image(推論結果を置く場所)
$ mkdir ml_stable_diffusion
$ cd ml_stable_diffusion
$ mkdir model
$ mkdir image
$ git clone https://github.com/apple/ml-stable-diffusion.git source

$ cd source
$ conda create -n pyenv38_ml_stable_diffusion python=3.8 -y
$ conda activate pyenv38_ml_stable_diffusion
$ pip install -e .

Hugging Face CLIを使用可能にする

apple/ml-stable-diffusionを動かすには、Hugging Faceから学習済みモデルを取得する必要があります。

ダウンロードして変換する処理は含まれていますが、事前にモデルをダウンロードするAPIを呼び出せる様にしておく必要があります。

# https://huggingface.co/docs/huggingface_hub/quick-start
$ pip install huggingface_jub
$ huggingface-cli login

モデルの変換処理

まずは同梱されている変換ライブラリを使ってモデルのダウンロードと変換を行います。

どのモデルを使用するかは–model-versionで指定します。指定しない場合は、 CompVis/stable-diffusion-v1-4 が指定されたものとみなされます。

ここではstabilityai/stable-diffusion-2-baseを指定し、出力先を../modelにしています。

使用している環境のメモリによっては個別生成をしないと変換途中で失敗する場合があります。

# 一回で生成する場合
$ python -m python_coreml_stable_diffusion.torch2coreml \
--convert-unet --convert-text-encoder --convert-vae-decoder \
--convert-safety-checker \
--model-version=stabilityai/stable-diffusion-2-base \
-o ../model

# MacBookAirのメモリが8GBytesタイプ等で足りない場合
$ python -m python_coreml_stable_diffusion.torch2coreml --convert-vae-decoder --model-version=stabilityai/stable-diffusion-2-base -o ../model
$ python -m python_coreml_stable_diffusion.torch2coreml --convert-unet --model-version=stabilityai/stable-diffusion-2-base -o ../model
$ python -m python_coreml_stable_diffusion.torch2coreml --convert-text-encoder --model-version=stabilityai/stable-diffusion-2-base -o ../model
$ python -m python_coreml_stable_diffusion.torch2coreml --convert-safety-checker --model-version=stabilityai/stable-diffusion-2-base -o ../model

画像の生成

生成したい画像の指示(いわゆる「呪文」)は、–promptで行います。

標準設定では–promptで指示した文字列がそのままフォルダ名になるため長い指示を行うとファイル名が生成出来ずに失敗してしまうため、試験動作目的の場合は短い文字列にしておきます。

ファイル名を指定可能にするにはpython_coreml_stable_diffusion/pipeline.pyに手をいれる必要があります。

$ python -m python_coreml_stable_diffusion.pipeline \
--prompt "a photo of an astronaut riding a horse on mars" \
--model-version=stabilityai/stable-diffusion-2-base \
-i ../model \
-o ../image --compute-unit CPU_AND_NE --seed 0