如何使用 Core ML 运行 Stable Diffusion

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

Core ML 模型可以利用 Apple 设备中所有可用的计算引擎：CPU、GPU 和 Apple Neural Engine（或 ANE，一种在 Apple Silicon Mac 和现代 iPhone/iPad 中可用的张量优化加速器）。根据模型及其运行的设备，Core ML 还可以混合和匹配计算引擎，例如，模型的某些部分可能在 CPU 上运行，而其他部分在 GPU 上运行。

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 基础和 2.1 基础
• coreml community 包括自定义微调模型
• 使用此 过滤器 返回所有可用的 Core ML 检查点

如果您找不到感兴趣的模型，我们建议您遵循 Apple 的 Converting Models to Core ML 说明。

选择要使用的 Core ML 变体

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

• 注意力类型 使用了n个块。注意力操作用于“关注”图像表示中不同区域之间的关系，并理解图像和文本表示如何相关。注意力的计算和内存消耗很大，因此存在不同的实现方式，以适应不同设备的硬件特性。对于Core ML Stable Diffusion模型，有两种注意力变体：

• split_einsum（由Apple引入）针对ANE设备进行了优化，这些设备在现代iPhone、iPad和M系列计算机中可用。
• “原始”注意力（在diffusers中使用的基础实现）仅与CPU/GPU兼容，不与ANE兼容。在CPU + GPU上使用original注意力运行模型可能比ANE更快。请参阅此性能基准以及社区提供的一些额外测量以获取更多细节。

• 支持的推理框架。

• packages适用于Python推理。这可用于在尝试将转换后的Core ML模型集成到原生应用程序之前进行测试，或者如果您想探索Core ML性能但不需要支持原生应用程序。例如，具有Web UI的应用程序完全可以使用Python Core ML后端。
• compiled模型是Swift代码所必需的。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"

mo
del_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

Core ML 在 Swift 中的推理

在 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 和潜在上采样器尚不可用。

Apple 的 转换和推理仓库 和我们自己的 swift-coreml-diffusers 仓库旨在作为技术演示，以帮助其他开发者在此基础上构建。

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

原生 Diffusers Swift 应用

一个简单的方法来在您自己的 Apple 硬件上运行 Stable Diffusion 是使用 我们的开源 Swift 仓库，它基于 diffusers 和 Apple 的转换和推理仓库。您可以研究代码，使用 Xcode 编译它，并根据您的需求进行适配。为了方便，App Store 中还有一个独立 Mac 应用，因此您无需处理代码或 IDE 即可使用它。如果您是开发者，并已确定 Core ML 是构建您的 Stable Diffusion 应用的最佳解决方案，那么您可以使用本指南的其余部分来开始您的项目。我们迫不及待想看看您会构建什么 🙂。