如何使用 Core ML 运行 Stable Diffusion

Core ML 是 Apple 框架支持的模型格式和机器学习库。如果您希望在 macOS 或 iOS/iPadOS 应用中运行 Stable Diffusion 模型，本指南将介绍如何将现有的 PyTorch 检查点转换为 Core ML 格式，并使用 Python 或 Swift 进行推理。

Core ML 模型可以利用 Apple 设备上所有可用的计算引擎：CPU、GPU 以及 Apple 神经网络引擎（即 ANE，一种张量优化的加速器，适用于 Apple Silicon Mac 以及现代 iPhone/iPad）。根据模型及其运行的设备，Core ML 还可以混合搭配不同的计算引擎，例如模型的一部分可能在 CPU 上运行，而另一部分则在 GPU 上运行。

您也可以在 Apple Silicon Mac 上使用 PyTorch 内置的 mps 加速器来运行 diffusers Python 代码库。此方法在 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 社区包含自定义微调模型
使用此 filter 可返回所有可用的 Core ML 检查点

如果您找不到感兴趣的模型，我们建议您按照 Apple 提供的将模型转换为 Core ML 说明进行操作。

选择要使用的 Core ML 变体

Stable Diffusion 模型可以转换为不同的 Core ML 变体，以用于不同的目的：

所使用的注意力块类型。注意力操作用于“关注”图像表示中不同区域之间的关系，并理解图像表示与文本表示之间的关联。注意力机制在计算和内存方面开销较大，因此存在不同的实现方式，以适应不同设备的硬件特性。对于 Core ML Stable Diffusion 模型，有两种注意力变体：
- split_einsum（由 Apple 引入）针对 ANE 设备进行了优化，该设备在现代 iPhone、iPad 和 M 系列计算机中可用。
- “原始”注意力机制（diffusers 中使用的基础实现）仅兼容 CPU/GPU，不支持 ANE。使用 original 注意力机制在 CPU + GPU 上运行模型可能比在 ANE 上运行更快。有关更多详细信息，请参阅此性能基准测试以及社区提供的额外测量数据。
支持的推理框架。
- packages 适用于 Python 推理。您可以在尝试将转换后的 Core ML 模型集成到原生应用之前，使用它来测试这些模型；或者，如果您希望探索 Core ML 的性能但不需要支持原生应用，也可以使用它。例如，具有 Web UI 的应用程序完全可以采用 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}")

推理

下载模型快照后，您可以使用 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 中运行推理比在 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}")

推理

要运行推理，请克隆 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，即 Stable Diffusion 默认使用的调度器，以及我们从 diffusers 实现中移植到 Swift 的 DPMSolverMultistepScheduler。我们建议使用 DPMSolverMultistepScheduler，因为它大约只需一半的步数即可达到相同的生成质量。
推理代码中支持负向提示词、无分类器引导尺度以及图生图任务。深度引导、ControlNet 和潜在上采样器等高级功能目前尚不可用。

苹果的转换与推理仓库以及我们自己的swift-coreml-diffusers仓库旨在作为技术演示，供其他开发者在此基础上进行开发。

如果您对某些缺失的功能有强烈需求，请随时提交功能请求，或者更好的方式是提交一个贡献 PR 🙂。

原生 Diffusers Swift 应用

在 Apple 硬件上运行 Stable Diffusion 的一个简便方法是使用我们的开源 Swift 仓库，该仓库基于 diffusers 以及 Apple 的转换与推理仓库构建。你可以研究代码，使用 Xcode 进行编译，并根据自身需求进行适配。为了方便起见，App Store 中还提供了一个独立的 Mac 应用，让你无需处理代码或 IDE 即可直接体验。如果你是开发者，并且已确定 Core ML 是构建 Stable Diffusion 应用的最佳方案，那么可以利用本指南的其余部分来启动你的项目。我们迫不及待想看看你会构建出什么 🙂。

Update on GitHub