Core MLでStable Diffusionを実行する方法

Core MLは、Appleのフレームワークがサポートするモデル形式および機械学習ライブラリです。macOSまたはiOS/iPadOSのアプリ内でStable Diffusionモデルを動かすことに興味がある場合、このガイドでは、既存のPyTorchチェックポイントをCore ML形式に変換し、PythonまたはSwiftで推論に使用する方法を説明します。

Core MLモデルは、Appleデバイスで利用可能なすべての計算エンジン（CPU、GPU、およびApple Neural Engine（ANE、Apple Silicon Macや最新のiPhone/iPadで利用可能なテンソル最適化アクセラレーター））を活用できます。モデルや実行環境となるデバイスによって、Core MLは計算エンジンを組み合わせることも可能です。例えば、モデルの一部はCPUで、別の一部はGPUで実行される場合があります。

diffusersのPythonコードベースは、PyTorchに組み込まれたmpsアクセラレーターを使用して、Apple Silicon Macでも実行できます。このアプローチについてはmpsガイドで詳しく説明されていますが、ネイティブアプリとは互換性がありません。

Stable Diffusion Core ML チェックポイント

Stable Diffusion の重み（またはチェックポイント）は PyTorch 形式で保存されているため、ネイティブアプリ内で使用するには、Core ML 形式に変換する必要があります。

幸いなことに、Appleのエンジニアはdiffusersをベースにした変換ツールを開発し、PyTorchのチェックポイントをCore MLに変換できるようにしました。

モデルの変換を行う前に、Hugging Face Hub を少し探索してみてください。興味のあるモデルが Core ML 形式で既に利用可能な場合があるためです：

Apple組織には、Stable Diffusionのバージョン1.4、1.5、2.0 base、および2.1 baseが含まれています。
coreml communityには、カスタムファインチューニングされたモデルが含まれています
このフィルターを使用して、利用可能なすべての Core ML チェックポイントを返します。

ご関心のあるモデルが見つからない場合は、Apple の Core ML へのモデル変換の手順に従うことをお勧めします。

使用するCore MLバリアントの選択

Stable Diffusionモデルは、異なる目的に使用されるさまざまなCore MLバリアントに変換できます。

使用されるアテンションブロックの種類です。アテンション操作は、画像表現内の異なる領域間の関係を「注目」し、画像表現とテキスト表現がどのように関連しているかを理解するために使用されます。アテンションは計算量とメモリ消費量が大きいため、異なるデバイスのハードウェア特性を考慮したさまざまな実装が存在します。Core ML Stable Diffusion モデルでは、2 つのアテンションバリアントがあります。
- split_einsum（Apple によって導入されました）は、最新の iPhone、iPad、および M シリーズのコンピュータで利用可能な ANE デバイス向けに最適化されています。
- 「original」アテンション（diffusersで使用されるベース実装）は、CPU/GPUとのみ互換性があり、ANEとは互換性がありません。originalアテンションを使用してCPU + GPUでモデルを実行する方が、ANEよりも高速になる場合があります。詳細については、このパフォーマンスベンチマークおよびコミュニティによって提供された追加の測定結果を参照してください。
サポートされている推論フレームワーク。
- packages は Python による推論に適しています。ネイティブアプリへの統合を試みる前に変換した Core ML モデルをテストする場合や、Core ML のパフォーマンスを調査したいがネイティブアプリのサポートは不要な場合に使用できます。例えば、Web UI を持つアプリケーションでは、Python の Core ML バックエンドを問題なく利用できます。
- Swiftコードにはcompiledモデルが必要です。Hub上のcompiledモデルでは、iOSおよびiPadOSデバイスとの互換性を確保するために、大きなUNetモデルの重みを複数のファイルに分割しています。これは--chunk-unet変換オプションに対応します。ネイティブアプリをサポートする場合は、compiledバリアントを選択する必要があります。

公式の Core ML Stable Diffusion モデルにはこれらのバリアントが含まれますが、コミュニティのものは異なる場合があります。

coreml-stable-diffusion-v1-4
├── README.md
├── original
│   ├── compiled
│   └── packages
└── split_einsum
    ├── compiled
    └── packages

必要に応じて、以下の方法でバリアントをダウンロードして使用できます。

PythonでのCore ML推論

PythonでCore ML推論を実行するには、以下のライブラリをインストールしてください:

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

モデルのチェックポイントをダウンロード

Pythonで推論を実行するには、compiled版はSwiftとのみ互換性があるため、packagesフォルダに保存されているバージョンのいずれかを使用してください。originalとsplit_einsumのどちらのattentionを使用するかを選択できます。

これは、original アテンションバリアントを Hub から models というディレクトリにダウンロードする方法です:

from huggingface_hub import snapshot_download
from pathlib import Path

repo_id = "apple/coreml-stable-diffusion-v1-4"
variant = "original/packages"

model_path = Path("./models") / (repo_id.split("/")[-1] + "_" + variant.replace("/", "_"))
snapshot_download(repo_id, allow_patterns=f"{variant}/*", local_dir=model_path, local_dir_use_symlinks=False)
print(f"Model downloaded at {model_path}")

推論

モデルのスナップショットをダウンロードした後、AppleのPythonスクリプトを使用してテストできます。

python -m python_coreml_stable_diffusion.pipeline --prompt "a photo of an astronaut riding a horse on mars" -i ./models/coreml-stable-diffusion-v1-4_original_packages/original/packages -o </path/to/output/image> --compute-unit CPU_AND_GPU --seed 93

ダウンロードしたチェックポイントのパスをスクリプトに-iフラグで渡します。--compute-unitは、推論に許可するハードウェアを指定します。以下のオプションのいずれかである必要があります: ALL、CPU_AND_GPU、CPU_ONLY、CPU_AND_NE。また、任意で出力パスと再現性のためのシードを指定できます。

推論スクリプトは、Stable Diffusion モデルのオリジナルバージョンである CompVis/stable-diffusion-v1-4 を使用していることを前提としています。別のモデルを使用する場合は、--model-version オプションを用いて、推論のコマンドラインでその Hub ID を指定する必要があります。これは、すでにサポートされているモデルや、ご自身で学習またはファインチューニングしたカスタムモデルにも適用できます。

例えば、stable-diffusion-v1-5/stable-diffusion-v1-5 を使用したい場合:

python -m python_coreml_stable_diffusion.pipeline --prompt "a photo of an astronaut riding a horse on mars" --compute-unit ALL -o output --seed 93 -i models/coreml-stable-diffusion-v1-5_original_packages --model-version stable-diffusion-v1-5/stable-diffusion-v1-5

SwiftでのCore ML推論

Swiftでの推論は、モデルがすでにmlmodelc形式でコンパイル済みのため、Pythonよりもわずかに高速です。これはモデルが読み込まれるアプリの起動時に顕著ですが、その後複数の生成を実行する場合は、体感できるほどの差にはなりません。

ダウンロード

Mac上でSwiftを使って推論を実行するには、compiledのチェックポイントバージョンのいずれかが必要です。前述の例と同様のPythonコードで、compiledのいずれかのバリアントをローカルにダウンロードすることをお勧めします。

from huggingface_hub import snapshot_download
from pathlib import Path

repo_id = "apple/coreml-stable-diffusion-v1-4"
variant = "original/compiled"

model_path = Path("./models") / (repo_id.split("/")[-1] + "_" + variant.replace("/", "_"))
snapshot_download(repo_id, allow_patterns=f"{variant}/*", local_dir=model_path, local_dir_use_symlinks=False)
print(f"Model downloaded at {model_path}")

推論

推論を実行するには、Appleのリポジトリをクローンしてください:

git clone https://github.com/apple/ml-stable-diffusion
cd ml-stable-diffusion

次に、AppleのコマンドラインツールであるSwift Package Managerを使用します:

swift run StableDiffusionSample --resource-path models/coreml-stable-diffusion-v1-4_original_compiled --compute-units all "a photo of an astronaut riding a horse on mars"

--resource-path には、前回のステップでダウンロードしたチェックポイントのいずれかを指定する必要があります。拡張子が .mlmodelc のコンパイル済み Core ML バンドルが含まれていることを確認してください。--compute-units には、all、cpuOnly、cpuAndGPU、cpuAndNeuralEngine のいずれかの値を指定する必要があります。

詳細については、Appleのリポジトリの指示を参照してください。

サポートされている Diffusers の機能

Core ML モデルおよび推論コードは、🧨 Diffusers の多くの機能、オプション、柔軟性をサポートしていません。以下は、考慮すべき主な制限事項です：

Core MLモデルは推論にのみ適しています。学習やファインチューニングには使用できません。
Swiftに移植されたスケジューラは2つのみです。Stable Diffusionで使用されるデフォルトのスケジューラと、diffusersの実装からSwiftに移植したDPMSolverMultistepSchedulerです。約半分のステップ数で同じ品質を生成できるため、DPMSolverMultistepSchedulerの使用を推奨します。
推論コードでは、ネガティブプロンプト、classifier-free guidance scale、image-to-imageタスクが利用可能です。depth guidance、ControlNet、latent upscalerなどの高度な機能は、まだ利用できません。

Appleの変換および推論リポジトリと、私たちが提供するswift-coreml-diffusersリポジトリは、他の開発者がこれらを基盤として開発を進められるよう、技術デモとして提供されています。

機能の欠落について強いご意見をお持ちの場合は、機能リクエストの作成、あるいは、より良い方法として、貢献PRの提出をぜひご検討ください 🙂。

ネイティブ Diffusers Swift アプリ

AppleのハードウェアでStable Diffusionを簡単に実行する1つの方法は、diffusersとAppleの変換・推論リポジトリをベースにしたオープンソースのSwiftリポジトリを使用することです。コードを調査し、Xcodeでコンパイルして、ご自身のニーズに合わせて調整できます。ご利便性のため、App StoreにスタンドアロンのMacアプリも提供されているため、コードやIDEを扱うことなく試すことができます。開発者であり、Stable Diffusionアプリの構築においてCore MLが最適なソリューションであると判断した場合は、このガイドの残りの部分を使ってプロジェクトを開始できます。皆さんが何を作るのか、今からとても楽しみです 🙂。

Update on GitHub