# 如何使用 Core ML 运行 Stable Diffusion

[Core ML](https://developer.apple.com/documentation/coreml) 是 Apple 框架支持的模型格式和机器学习库。如果您希望在 macOS 或 iOS/iPadOS 应用中运行 Stable Diffusion 模型，本指南将介绍如何将现有的 PyTorch 检查点转换为 Core ML 格式，并使用 Python 或 Swift 进行推理。

Core ML 模型可以利用 Apple 设备上所有可用的计算引擎：CPU、GPU 以及 Apple 神经网络引擎（即 ANE，一种张量优化的加速器，适用于 Apple Silicon Mac 和现代 iPhone/iPad）。根据模型及其运行的设备，Core ML 还可以混合搭配不同的计算引擎，例如模型的一部分可能在 CPU 上运行，而另一部分则在 GPU 上运行。

> [!TIP]
> 您也可以在 Apple Silicon Mac 上使用 PyTorch 内置的 `mps` 加速器来运行 `diffusers` Python 代码库。此方法在 [mps 指南](mps) 中有详细说明，但它与原生应用不兼容。

## Stable Diffusion Core ML 检查点

Stable Diffusion 的权重（或检查点）以 PyTorch 格式存储，因此在使用它们于原生应用之前，需要将其转换为 Core ML 格式。

幸运的是，Apple 工程师基于 `diffusers` 开发了[一个转换工具](https://github.com/apple/ml-stable-diffusion#-converting-models-to-core-ml)，用于将 PyTorch 检查点转换为 Core ML。

在转换模型之前，请先花点时间浏览 Hugging Face Hub——您感兴趣的模型很可能已经以 Core ML 格式提供：

- [Apple](https://huggingface.co/apple) 组织包含 Stable Diffusion 1.4、1.5、2.0 base 和 2.1 base 版本
- [coreml 社区](https://huggingface.co/coreml-community)包含自定义微调模型
- 使用此 [filter](https://huggingface.co/models?pipeline_tag=text-to-image&library=coreml&p=2&sort=likes) 可返回所有可用的 Core ML 检查点

如果您找不到感兴趣的模型，我们建议您按照 Apple 提供的 [将模型转换为 Core ML](https://github.com/apple/ml-stable-diffusion#-converting-models-to-core-ml) 说明进行操作。

## 选择要使用的 Core ML 变体

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

- 所使用的注意力块类型。注意力操作用于“关注”图像表示中不同区域之间的关系，并理解图像表示与文本表示之间的关联。注意力计算在计算和内存方面开销较大，因此存在不同的实现方式，以适应不同设备的硬件特性。对于 Core ML Stable Diffusion 模型，有两种注意力变体：
    * `split_einsum`（[由 Apple 引入](https://machinelearning.apple.com/research/neural-engine-transformers)）针对 ANE 设备进行了优化，该设备在现代 iPhone、iPad 和 M 系列计算机中可用。
    * “原始”注意力机制（`diffusers` 中使用的基础实现）仅兼容 CPU/GPU，不支持 ANE。使用 `original` 注意力机制在 CPU + GPU 上运行模型可能比在 ANE 上运行*更快*。有关更多详细信息，请参阅[此性能基准测试](https://huggingface.co/blog/fast-mac-diffusers#performance-benchmarks)以及社区提供的[额外测量数据](https://github.com/huggingface/swift-coreml-diffusers/issues/31)。

- 支持的推理框架。
    * `packages` 适用于 Python 推理。您可以在尝试将转换后的 Core ML 模型集成到原生应用之前，使用它来测试这些模型；或者，如果您希望探索 Core ML 的性能但不需要支持原生应用，也可以使用它。例如，具有 Web UI 的应用程序完全可以采用 Python Core ML 后端。
    * Swift 代码需要 `compiled` 模型。Hub 中的 `compiled` 模型将大型 UNet 模型权重拆分为多个文件，以兼容 iOS 和 iPadOS 设备。这对应于 [`--chunk-unet` 转换选项](https://github.com/apple/ml-stable-diffusion#-converting-models-to-core-ml)。如果您希望支持原生应用，则需要选择 `compiled` 变体。

官方 Core ML Stable Diffusion [模型](https://huggingface.co/apple/coreml-stable-diffusion-v1-4/tree/main)包含以下变体，但社区版本可能有所不同：

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

您可以按照以下说明下载并使用所需的变体。

## Python 中的 Core ML 推理

安装以下库以在 Python 中运行 Core ML 推理：

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

### 下载模型检查点

要在 Python 中运行推理，请使用 `packages` 文件夹中存储的版本，因为 `compiled` 版本仅与 Swift 兼容。您可以选择使用 `original` 或 `split_einsum` 注意力机制。

这是如何从 Hub 下载 `original` 注意力变体到名为 `models` 的目录：

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

### 推理[[python-inference]]

下载模型快照后，您可以使用 Apple 提供的 Python 脚本对其进行测试。

```shell
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`](https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5)：

```shell
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
```

## Swift 中的 Core ML 推理

在 Swift 中运行推理比在 Python 中稍快，因为模型已经以 `mlmodelc` 格式编译。在应用启动加载模型时，这种速度差异会比较明显，但如果你随后运行多次生成，则不太容易察觉。

### 下载

要在 Mac 上使用 Swift 运行推理，您需要一个`compiled`检查点版本。我们建议您使用与前面示例类似的 Python 代码在本地下载它们，但需选择`compiled`变体之一：

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

### 推理[[swift-inference]]

要运行推理，请克隆 Apple 的仓库：

```bash
git clone https://github.com/apple/ml-stable-diffusion
cd ml-stable-diffusion
```

然后使用 Apple 的命令行工具 [Swift Package Manager](https://www.swift.org/package-manager/#)：

```bash
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 仓库中的说明](https://github.com/apple/ml-stable-diffusion)。

## 支持的 Diffusers 功能

Core ML 模型和推理代码不支持 🧨 Diffusers 的许多功能、选项和灵活性。以下是需要注意的一些限制：

- Core ML 模型仅适用于推理，不能用于训练或微调。
- 目前只有两个调度器被移植到了 Swift，即 Stable Diffusion 默认使用的调度器，以及我们从 `diffusers` 实现中移植到 Swift 的 `DPMSolverMultistepScheduler`。我们建议使用 `DPMSolverMultistepScheduler`，因为它大约只需一半的步数即可达到相同的生成质量。
- 推理代码中支持负向提示词、无分类器引导尺度以及图生图任务。深度引导、ControlNet 和潜在上采样器等高级功能目前尚不可用。

苹果的[转换与推理仓库](https://github.com/apple/ml-stable-diffusion)以及我们自己的[swift-coreml-diffusers](https://github.com/huggingface/swift-coreml-diffusers)仓库旨在作为技术演示，供其他开发者在此基础上进行开发。

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

## 原生 Diffusers Swift 应用

在 Apple 硬件上运行 Stable Diffusion 的一个简便方法是使用[我们的开源 Swift 仓库](https://github.com/huggingface/swift-coreml-diffusers)，该仓库基于 `diffusers` 以及 Apple 的转换与推理仓库构建。你可以研究代码，使用 [Xcode](https://developer.apple.com/xcode/) 进行编译，并根据自身需求进行适配。为了方便起见，App Store 中还提供了一个[独立的 Mac 应用](https://apps.apple.com/app/diffusers/id1666309574)，让你无需处理代码或 IDE 即可直接体验。如果你是开发者，并且已确定 Core ML 是构建 Stable Diffusion 应用的最佳方案，那么可以利用本指南的其余部分来启动你的项目。我们迫不及待想看看你会构建出什么 🙂。