如何在 Apple Silicon 上用 Core ML 运行 Stable Diffusion

多亏了 Apple 工程师，你现在可以使用 Core ML 在 Apple Silicon 上运行 Stable Diffusion！

这个 Apple 仓库提供了基于 🧨 Diffusers 的转换脚本和推理代码，我们非常喜欢！为了让你尽可能轻松上手，我们自己转换了权重，并将模型的 Core ML 版本放到了 Hugging Face Hub 上。

更新：在这篇文章发布几周后，我们创建了一个原生 Swift 应用，你可以用它轻松在自己的硬件上运行 Stable Diffusion。我们发布了 Mac App Store 上的应用以及源代码，以便其他项目使用。

本文的剩余部分将指导你如何在自己的代码中使用转换后的权重，或者自己转换额外的权重。

可用的检查点

官方的 Stable Diffusion 检查点已经转换好，可以直接使用：

Stable Diffusion v1.4：已转换原始
Stable Diffusion v1.5：已转换原始
Stable Diffusion v2 base：已转换原始
Stable Diffusion v2.1 base：已转换原始

Core ML 支持你设备中所有可用的计算单元：CPU、GPU 和 Apple 的神经引擎（NE）。Core ML 还可以在不同设备上运行模型的不同部分，以最大化性能。

每个模型都有多个变体，根据你使用的硬件，性能可能不同。我们建议你尝试它们，并坚持使用在你的系统上表现最好的那个。继续阅读了解详情。

性能说明

每个模型有多个变体：

"Original" 注意力与 "split_einsum"。这是关键注意力块的两种替代实现。split_einsum 是之前由 Apple 引入的，并且兼容所有计算单元（CPU、GPU 和 Apple 的神经引擎）。另一方面，original 只兼容 CPU 和 GPU。不过，在某些设备上 original 可能比 split_einsum 更快，所以一定要试试！
"ML Packages" 与 "Compiled" 模型。前者适用于 Python 推理，而 compiled 版本是 Swift 代码所需的。Hub 中的 compiled 模型将大型 UNet 模型权重拆分成多个文件，以兼容 iOS 和 iPadOS 设备。这对应于 --chunk-unet 转换选项。

在撰写本文时，我们在我的 MacBook Pro（M1 Max，32 个 GPU 核心，64 GB）上使用以下组合获得了最佳结果：

original 注意力。
all 计算单元（详情见下一节）。
macOS Ventura 13.1 Beta 4 (22C5059b)。

使用这些，用 Core ML 版本的 Stable Diffusion v1.4 生成一张图像需要 18 秒 🤯。

⚠️ 注意

macOS Ventura 13.1 中引入了多项 Core ML 改进，Apple 的实现需要这些改进。如果你使用旧版本的 macOS，可能会得到黑色图像——并且速度慢得多。

每个模型仓库都组织成树状结构，提供这些不同的变体：

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

你可以下载并使用你需要的变体，如下所示。

Python 中的 Core ML 推理

先决条件

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

下载模型检查点

要在 Python 中运行推理，你必须使用存储在 packages 文件夹中的某个版本，因为编译后的版本仅与 Swift 兼容。你可以选择使用 original 或 split_einsum 注意力风格。

以下是从 Hub 下载 original 注意力变体的方法：

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

上面的代码会将下载的模型快照放在名为 models 的目录中。

推理

下载模型快照后，运行推理最简单的方法是使用 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 -o </path/to/output/image> --compute-unit ALL --seed 93

<output-mlpackages-directory> 应指向你在上一步中下载的检查点，--compute-unit 表示你希望允许用于推理的硬件。它必须是以下选项之一：ALL、CPU_AND_GPU、CPU_ONLY、CPU_AND_NE。你还可以提供可选的输出路径，以及用于可复现性的种子。

推理脚本假定使用 Stable Diffusion 模型的原始版本，在 Hub 中存储为 CompVis/stable-diffusion-v1-4。如果你使用其他模型，你必须在推理命令行中使用 --model-version 选项指定其 Hub id。这既适用于已支持的模型，也适用于你自己训练或微调的自定义模型。

对于 Stable Diffusion 1.5（Hub id：runwayml/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 runwayml/stable-diffusion-v1-5

对于 Stable Diffusion 2 base（Hub id：stabilityai/stable-diffusion-2-base）：

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-2-base_original_packages --model-version stabilityai/stable-diffusion-2-base

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

然后使用 Swift Package Manager 的功能来使用 Apple 的命令行工具：

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 仓库中的说明。

自带模型

如果你创建了与 Stable Diffusion 兼容的自有模型（例如，如果你使用了 Dreambooth、Textual Inversion 或微调），那么你必须自行转换模型。幸运的是，Apple 提供了一个转换脚本，允许你这样做。

对于此任务，我们建议你遵循这些说明。

后续步骤

我们对这带来的机遇感到非常兴奋，迫不及待地想看到社区能在此基础上创造出什么。一些潜在的想法包括：

适用于 Mac、iPhone 和 iPad 的原生高质量应用。
为 Swift 引入更多调度器，以实现更快的推理。
更多管道和任务。
探索量化技术和进一步优化。

期待看到你的创作！