感谢苹果工程师,现在您可以使用 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 基础版:转换的原始版本
- Stable Diffusion v2.1 基础版:转换的原始版本
Core ML 支持设备中可用的所有计算单元:CPU、GPU 和苹果的神经引擎 (NE)。Core ML 还可以在不同的设备上运行模型的不同部分,以最大化性能。
每个模型都有几个变体,根据您使用的硬件不同,可能会产生不同的性能。我们建议您尝试它们,并选择在您的系统中效果最好的一个。继续阅读以获取详细信息。
性能说明
每个模型有几个变体:
- “Original” attention 和 “split_einsum”。这是关键的注意力块的两种替代实现。
split_einsum
是之前由 Apple 引入的,与所有计算单元 (CPU、GPU 和苹果的神经引擎) 兼容。另一方面,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
attention。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
attention variant 的方法:
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_path}")
以上代码将把下载的模型快照放在名为models
的目录中。
推理
一旦你下载了模型的快照,最简单的运行推理的方法是使用苹果的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 <输出图片的路径> --compute-unit ALL --seed 93
<output-mlpackages-directory>
应该指向你在上一步下载的检查点,并且--compute-unit
表示你想要用于推理的硬件。它必须是以下选项之一:ALL
、CPU_AND_GPU
、CPU_ONLY
、CPU_AND_NE
。你也可以提供一个可选的输出路径和一个用于重现的种子。
推理脚本假设你使用的是存储在Hub中的Stable Diffusion模型的原始版本,其Hub id为CompVis/stable-diffusion-v1-4
。如果你使用其他模型,你必须在推理命令行中指定其Hub id,使用--model-version
选项。这对于已经支持的模型和你自己训练或微调的自定义模型都适用。
对于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_path}")
推理
要运行推理,请克隆苹果的repo:
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 "一张火星上骑马的宇航员照片"
您需要在--resource-path
中指定先前步骤中下载的检查点之一,请确保它包含扩展名为.mlmodelc
的编译后的Core ML模型包。而--compute-units
必须是以下值之一:all
,cpuOnly
,cpuAndGPU
,cpuAndNeuralEngine
。
更多详细信息,请参考Apple的存储库中的说明。
使用自己的模型
如果您已经创建了与Stable Diffusion兼容的自己的模型(例如,如果您使用了Dreambooth、Textual Inversion或微调),那么您需要自己转换这些模型。幸运的是,Apple提供了一个转换脚本供您使用。
对于这个任务,我们建议您按照这些说明进行操作。
下一步
我们对这些机会感到非常兴奋,迫不及待地想看到社区能从中创造出什么。一些潜在的创意包括:
- 为Mac、iPhone和iPad开发本地高质量应用。
- 为Swift带来更多调度器,以实现更快的推断。
- 添加其他的流水线和任务。
- 探索量化技术和进一步的优化。
期待看到您的创作!