⚒️

Apple Sillicon環境で、Stable DiffusionのCoreMLモデルを動かす

2022/12/03に公開

2022/12/1、AppleからStable DiffusionのCoreML対応モデルが公開されました。合わせて、Hugging Faceの公式ブログ[1]にてセットアップ方法も書かれました。そこで、Apple Sillicon環境でモデルダウンロード + プロンプトから画像作成までをPythonにて実施したので共有いたします。CoreMLモデルを使うことで、Neural Engineを最大限活用しながら、推論処理を行えます。

前置き

  • 本稿のサンプルコードを https://github.com/Niccari/coreml-stable-diffusion-cli-example においているので、必要に応じご利用ください。
  • モデルのダウンロードおよび推論実行には、Hugging FaceのWrite権限つきアクセストークンが必要です。本稿ではアカウント作成・アクセストークンの発行方法については割愛すること、ご了承ください。

1. 必要ライブラリのインストール

huggingface_hub および ml-stable-diffusion をインストールします。

$ pip install huggingface_hub
$ pip install git+https://github.com/apple/ml-stable-diffusion

なお、pipenvを使っている場合は、

$pipenv install huggingface_hub git+https://github.com/apple/ml-stable-diffusion#egg=python-coreml-stable-diffusion==0.1.0

という形でパッケージ名とバージョンを明示する必要があります。

2. Hugging Faceへのログイン

以下コマンドにより、Hugging Faceにログインします(ログイン済みであれば本ステップは不要です)。アクセストークンを求められるので、Write権限のあるアクセストークンを指定します。

$ 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) n
Token is valid.
Your token has been saved to /Users/niccari/.huggingface/token
Login successful

なお、アクセストークンが間違っていたり、read権限しかない場合はValueError: Invalid token passed!が発生します。

3. モデルのダウンロード

Appleからは、以下のモデルが提供されています。

そこで、公式ブログの内容をもとに、以下コードなどで対話形式でダウンロードできるようにしてみます。ターミナルにて、

  • 2を入力する -> v1.5をダウンロード
  • 3を入力する -> 2-baseをダウンロード
  • その他 -> v1.4をダウンロード

となるようになっています。

from huggingface_hub import snapshot_download
from huggingface_hub.file_download import repo_folder_name
from pathlib import Path
import shutil


def download_model(repo_id, variant, output_dir):
    destination = Path(output_dir) / (repo_id.split("/")[-1] + "_" + variant.replace("/", "_"))
    if destination.exists():
        raise FileExistsError(f"Model already exists at {destination}")
    
    downloaded = snapshot_download(
        repo_id, allow_patterns=f"{variant}/*", cache_dir=output_dir
    )
    downloaded_bundle = Path(downloaded) / variant
    shutil.copytree(downloaded_bundle, destination)

    cache_folder = Path(output_dir) / repo_folder_name(repo_id=repo_id, repo_type="model")
    shutil.rmtree(cache_folder)
    return destination


def to_version(input_text: str):
    if input_text == "3":
        return "apple/coreml-stable-diffusion-2-base"
    elif input_text == "2":
        return "apple/coreml-stable-diffusion-v1-5"
    else:
        return "apple/coreml-stable-diffusion-v1-4"


if __name__ == "__main__":
    print("Note: you must log in to Hugging face hub.")
    input_text = input(
        "Model version? - [1,*] v1.4, [2] v1.5, [3] v2.0 > ")
    repo_id = to_version(input_text)

    model_path = download_model(
        repo_id, "original/packages", output_dir="./models")
    print(f"Model downloaded at {model_path}")

こちらを実行すると、以下の様にモデルのダウンロードが行われます。ログの最後に、ダウンロードしたモデルの保存先を表示します。

実行ログの例

Note: you must log in to Hugging face hub.
Model version? - [1,*] v1.4, [2] v1.5, [3] v2.0 > 3
Downloading: 100%|████████████████████████████████████████████████████████████████████| 832M/832M [00:39<00:00, 21.1MB/s]
Downloading: 100%|███████████████████████████████████████████████████████████████████████| 617/617 [00:00<00:00, 213kB/s]
Downloading: 100%|█████████████████████████████████████████████████████████████████████| 147k/147k [00:00<00:00, 432kB/s]
Downloading: 100%|██████████████████████████████████████████████████████████████████| 99.0M/99.0M [00:04<00:00, 21.3MB/s]
Downloading: 100%|███████████████████████████████████████████████████████████████████████| 617/617 [00:00<00:00, 488kB/s]
Fetching 18 files: 100%|█████████████████████████████████████████████████████████████████| 18/18 [00:52<00:00,  2.92s/it]
Model downloaded at models/coreml-stable-diffusion-v1-5_original_packages

3. 推論の実行

以下のコマンドにより、ダウンロードしたモデルとプロンプトを指定して実行します。-iオプションでは使いたいモデルのパス、--promptオプションでプロンプトとなる文を指定します。

$ python -m python_coreml_stable_diffusion.pipeline \
  --prompt "succurents grow in clusters on a desert" \
  -i models/coreml-stable-diffusion-v1-4_original_packages \ 
  -o outputs/ \
  --compute-unit ALL \
  --seed 93

すると、-oオプションで指定したフォルダに結果画像が出力されます。MacBook Air M1 2020(メモリ8GB)だと、30分弱かかりました。

succurents grow in clusters on a desert (v1.4)

なお、v1.5以上のモデルを使う場合はさらに--model-versionオプションが必要です。具体的には、

です。これらの指定がない場合、デフォルトのCompVis/stable-diffusion-v1-4がモデル名として扱われ、FileNotFoundError: text_encoder CoreML model doesn't exist at models/coreml-stable-diffusion-v1-5_original_packages/Stable_Diffusion_version_CompVis_stable-diffusion-v1-4_text_encoder.mlpackageでエラーになります。

以上です!

脚注
  1. https://huggingface.co/blog/diffusers-coreml ↩︎

  2. 別途認可が必要です。認可されていない場合、Access to model runwayml/stable-diffusion-v1-5 is restricted and you are not in the authorized list. Visit https://huggingface.co/runwayml/stable-diffusion-v1-5 to ask for access.となります。 ↩︎

Discussion