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

感谢 Apple 工程师的努力，现在你可以使用 Core ML 在 Apple Silicon 上运行 Stable Diffusion！

这个 Apple 仓库提供了基于 🧨 Diffusers 的转换脚本和推理代码，我们非常喜欢！为了让你尽可能轻松上手，我们自行转换了权重，并将模型的 Core ML 版本放到了 HF Mirror 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 也可以在不同设备上运行模型的不同部分，以实现最佳性能。

每种模型都有多个变体，根据你使用的硬件，性能表现可能不同。我们建议你逐一尝试，并选择最适合你系统的版本。继续阅读以了解详情。

性能说明

每种模型有多个变体：

“原始”注意力与“split_einsum”。这是关键注意力块的两种替代实现。split_einsum 是由 Apple 先前引入的，兼容所有计算单元（CPU、GPU 和 Apple 的神经引擎）。而 original 仅兼容 CPU 和 GPU。尽管如此，在某些设备上 original 可能比 split_einsum 更快，所以务必试试！
“ML Packages”与“Compiled”模型。前者适用于 Python 推理，而 Swift 代码则需要 compiled 版本。Hub 中的 compiled 模型将大型 UNet 模型权重拆分为多个文件，以兼容 iOS 和 iPadOS 设备。这对应于 --chunk-unet 转换选项。

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

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

使用这些配置，用 Stable Diffusion v1.4 的 Core ML 版本生成一张图像仅需 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 引入更多调度器，以实现更快的推理速度。
更多的管道和任务。
探索量化技术和进一步的优化。

期待看到你的创作！