如何使用 Core ML 執行 Stable Diffusion

Core ML 是蘋果框架支援的模型格式和機器學習庫。如果您有興趣在 macOS 或 iOS/iPadOS 應用程式中執行 Stable Diffusion 模型，本指南將向您展示如何將現有的 PyTorch 檢查點轉換為 Core ML 格式，並使用它們進行 Python 或 Swift 推理。

Core ML 模型可以利用蘋果裝置中所有可用的計算引擎：CPU、GPU 和蘋果神經引擎（ANE，一種在蘋果晶片 Mac 和現代 iPhone/iPad 中可用的張量最佳化加速器）。根據模型和執行裝置的具體情況，Core ML 還可以混合搭配使用計算引擎，例如，模型的一部分可以在 CPU 上執行，而另一部分則可以在 GPU 上執行。

Stable Diffusion Core ML 檢查點

Stable Diffusion 的權重（或檢查點）以 PyTorch 格式儲存，因此在使用它們構建原生應用之前，您需要將它們轉換為 Core ML 格式。

值得慶幸的是，蘋果工程師開發了一個基於 `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 檢查點

如果您找不到感興趣的模型，我們建議您遵循蘋果的將模型轉換為 Core ML 的說明。

選擇要使用的 Core ML 變體

Stable Diffusion 模型可以轉換為不同的 Core ML 變體，以適應不同目的

• 使用的注意力塊型別。注意力操作用於“關注”影像表示中不同區域之間的關係，並理解影像和文字表示如何關聯。注意力機制是計算和記憶體密集型的，因此存在不同的實現，考慮了不同裝置的硬體特性。對於 Core ML Stable Diffusion 模型，有兩種注意力變體

  • `split_einsum`（由蘋果引入）針對 ANE 裝置進行了最佳化，這些裝置在現代 iPhone、iPad 和 M 系列電腦中可用。
  • “原始”注意力（`diffusers` 中使用的基礎實現）僅與 CPU/GPU 相容，而與 ANE 不相容。在 CPU + GPU 上使用 `original` 注意力執行模型可能比在 ANE 上*更快*。有關更多詳細資訊，請參閱此效能基準測試以及社群提供的一些附加測量資料。

• 支援的推理框架。

  • `packages` 適用於 Python 推理。這可以在嘗試將轉換後的 Core ML 模型整合到原生應用之前進行測試，或者如果您想探索 Core ML 效能但不需要支援原生應用。例如，一個帶有網頁介面的應用程式可以完美地使用 Python Core ML 後端。
  • Swift 程式碼需要 `compiled` 模型。Hub 中的 `compiled` 模型將大型 UNet 模型權重分割成多個檔案，以相容 iOS 和 iPadOS 裝置。這對應於`--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 中進行推理，請使用 `packages` 資料夾中儲存的版本之一，因為 `compiled` 版本僅與 Swift 相容。您可以選擇使用 `original` 或 `split_einsum` 注意力。

以下是如何從 Hub 下載 `original` 注意力變體到名為 `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}")

推理

下載了模型的快照後，您可以使用蘋果的 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 中進行推理比在 Python 中稍快，因為模型已經編譯成 `mlmodelc` 格式。這在應用啟動載入模型時很明顯，但在之後執行多次生成時則不明顯。

下載

要在 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}")

推理

要執行推理，請克隆蘋果的倉庫

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

然後使用蘋果的命令列工具，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`。

有關更多詳細資訊，請參閱蘋果倉庫中的說明。

支援的 Diffusers 功能

Core ML 模型和推理程式碼不支援 🧨 Diffusers 的許多功能、選項和靈活性。以下是一些需要注意的限制

• Core ML 模型僅適用於推理。它們不能用於訓練或微調。
• 只有兩個排程器被移植到了 Swift，即 Stable Diffusion 使用的預設排程器和 `DPMSolverMultistepScheduler`，我們從我們的 `diffusers` 實現中將其移植到了 Swift。我們建議您使用 `DPMSolverMultistepScheduler`，因為它能在大約一半的步數內產生相同的質量。
• 推理程式碼中提供了負向提示、無分類器指導縮放和影像到影像任務。高階功能，如深度指導、ControlNet 和潛在空間升頻器尚不可用。

蘋果的轉換和推理倉庫以及我們自己的swift-coreml-diffusers倉庫旨在作為技術演示，以便其他開發者可以在此基礎上進行構建。

如果您對任何缺失的功能有強烈的需求，請隨時提交功能請求，或者更好的是，提交貢獻 PR 🙂。

原生 Diffusers Swift 應用

在您自己的蘋果硬體上執行 Stable Diffusion 的一個簡單方法是使用我們基於 `diffusers` 和蘋果的轉換與推理倉庫的開源 Swift 倉庫。您可以研究程式碼，用 Xcode 編譯它，並根據自己的需求進行調整。為方便起見，App Store 中還有一個獨立的 Mac 應用，這樣您就可以直接體驗，而無需處理程式碼或 IDE。如果您是一名開發者，並確定 Core ML 是構建 Stable Diffusion 應用的最佳解決方案，那麼您可以使用本指南的其餘部分來開始您的專案。我們迫不及待地想看到您將構建出什麼 🙂。