在 Apple Silicon 上通过 Core ML 使用 Stable Diffusion
感谢苹果公司的工程师们，你现在可以在 Apple Silicon 上通过 Core ML 运行 Stable Diffusion 了！
这个 Apple 的 repo 提供了基于 🧨 Diffusers 的转换脚本和推理代码，我们非常喜欢！为了让你尽可能轻松地使用，我们自己转换了权重，并将模型的 Core ML 版本放在了 Hugging Face Hub 上。
更新：在这篇文章发布几周后，我们创建了一个原生的 Swift 应用，你可以用它在自己的硬件上轻松运行 Stable Diffusion。我们在 Mac App Store 上发布了一个应用，并开源了源代码，以允许其他项目使用它。
本文的其余部分将指导你如何在自己的代码中使用转换后的权重，或者自己转换其他权重。
可用的 checkpoints
官方的 Stable Diffusion checkpoints 已经转换完毕并可供使用
- Stable Diffusion v1.4: 转换版 原始版
- Stable Diffusion v1.5: 转换版 原始版
- Stable Diffusion v2 base: 转换版 原始版
- Stable Diffusion v2.1 base: 转换版 原始版
Core ML 支持你设备上所有可用的计算单元：CPU、GPU 和苹果的神经网络引擎 (NE)。Core ML 还可以将模型的不同部分在不同设备上运行以最大化性能。
每个模型都有几个变体，根据你使用的硬件可能会产生不同的性能。我们建议你尝试一下，并选择在你的系统上效果最好的那个。详情请继续阅读。
性能说明
每个模型有几个变体
- “Original” attention vs “split_einsum”。这是关键的 attention 模块的两种替代实现。
split_einsum是由苹果之前引入的，并且与所有计算单元 (CPU、GPU 和苹果的神经网络引擎) 兼容。而original只与 CPU 和 GPU 兼容。尽管如此，original在某些设备上可能比split_einsum更快，所以一定要试试！ - “ML Packages” vs “Compiled” 模型。前者适用于 Python 推理，而
compiled版本是 Swift 代码所必需的。Hub 中的compiled模型将大型 UNet 模型权重分割成多个文件，以兼容 iOS 和 iPadOS 设备。这对应于--chunk-unet转换选项。
在撰写本文时，我们在我的 MacBook Pro (M1 Max, 32 核 GPU, 64 GB) 上使用以下组合获得了最佳结果
originalattention。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
下载模型 Checkpoints
要在 Python 中运行推理，你必须使用存储在 packages 文件夹中的版本之一，因为编译后的版本只与 Swift 兼容。你可以选择使用 original 还是 split_einsum 的 attention 风格。
以下是如何从 Hub 下载 original attention 变体的方法
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> 应该指向你在上一步下载的 checkpoint，而 --compute-unit 则指明你希望用于推理的硬件。它必须是以下选项之一：ALL、CPU_AND_GPU、CPU_ONLY、CPU_AND_NE。你也可以提供一个可选的输出路径，以及一个用于可复现性的种子 (seed)。
推理脚本假设使用的是 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 版本的 checkpoint。我们建议你使用类似于我们上面展示的 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 中指定上一步下载的 checkpoint 之一，所以请确保它包含扩展名为 .mlmodelc 的已编译 Core ML 包。--compute-units 的值必须是以下之一：all、cpuOnly、cpuAndGPU、cpuAndNeuralEngine。
更多详情，请参考Apple 仓库中的说明。
使用你自己的模型
如果你创建了自己与 Stable Diffusion 兼容的模型 (例如，如果你使用了 Dreambooth、Textual Inversion 或微调)，那么你必须自己转换模型。幸运的是，Apple 提供了一个转换脚本，可以让你做到这一点。
对于这项任务，我们建议你遵循这些说明。
下一步
我们对这带来的机会感到非常兴奋，并迫不及待地想看看社区能从中创造出什么。一些潜在的想法是
- 适用于 Mac、iPhone 和 iPad 的原生、高质量应用。
- 将更多的调度器引入 Swift，以实现更快的推理。
- 额外的流水线和任务。
- 探索量化技术和进一步的优化。
期待看到你的创作！
