如何使用 Core ML 运行 Stable Diffusion

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

Core ML 模型可以利用苹果设备中所有可用的计算引擎：CPU、GPU 和苹果神经引擎（ANE，一种在苹果芯片 Mac 和现代 iPhone/iPad 中可用的张量优化加速器）。根据模型和运行设备的具体情况，Core ML 还可以混合搭配使用计算引擎，例如，模型的一部分可以在 CPU 上运行，而另一部分则可以在 GPU 上运行。

Stable Diffusion Core ML 检查点

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

值得庆幸的是，苹果工程师开发了一个基于 `diffusers` 的转换工具，可以将 PyTorch 检查点转换为 Core ML。

不过，在转换模型之前，请花点时间浏览 Hugging Face Hub——您感兴趣的模型很可能已经有 Core ML 格式的版本了。

• Apple 组织包含了 Stable Diffusion 1.4、1.5、2.0 base 和 2.1 base 版本
• coreml community 包含了自定义微调模型
• 使用此筛选器返回所有可用的 Core ML 检查点

如果您找不到感兴趣的模型，我们建议您遵循苹果的将模型转换为 Core ML 的说明。

选择要使用的 Core ML 变体

Stable Diffusion 模型可以转换为不同的 Core ML 变体，以适应不同目的

• 使用的注意力块类型。注意力操作用于“关注”图像表示中不同区域之间的关系，并理解图像和文本表示如何关联。注意力机制是计算和内存密集型的，因此存在不同的实现，考虑了不同设备的硬件特性。对于 Core ML Stable Diffusion 模型，有两种注意力变体

  • `split_einsum`（由苹果引入）针对 ANE 设备进行了优化，这些设备在现代 iPhone、iPad 和 M 系列电脑中可用。
  • “原始”注意力（`diffusers` 中使用的基础实现）仅与 CPU/GPU 兼容，而与 ANE 不兼容。在 CPU + GPU 上使用 `original` 注意力运行模型可能比在 ANE 上*更快*。有关更多详细信息，请参阅此性能基准测试以及社区提供的一些附加测量数据。

• 支持的推理框架。

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

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 脚本进行测试。

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

在 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}")

推理

要运行推理，请克隆苹果的仓库

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 "a photo of an astronaut riding a horse on mars"

您必须在 `--resource-path` 中指定在上一步下载的检查点之一，所以请确保它包含扩展名为 `.mlmodelc` 的已编译 Core ML 包。`--compute-units` 必须是以下值之一：`all`、`cpuOnly`、`cpuAndGPU`、`cpuAndNeuralEngine`。

有关更多详细信息，请参阅苹果仓库中的说明。

支持的 Diffusers 功能

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

• Core ML 模型仅适用于推理。它们不能用于训练或微调。
• 只有两个调度器被移植到了 Swift，即 Stable Diffusion 使用的默认调度器和 `DPMSolverMultistepScheduler`，我们从我们的 `diffusers` 实现中将其移植到了 Swift。我们建议您使用 `DPMSolverMultistepScheduler`，因为它能在大约一半的步数内产生相同的质量。
• 推理代码中提供了负向提示、无分类器指导缩放和图像到图像任务。高级功能，如深度指导、ControlNet 和潜在空间升频器尚不可用。

苹果的转换和推理仓库以及我们自己的swift-coreml-diffusers仓库旨在作为技术演示，以便其他开发者可以在此基础上进行构建。

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

原生 Diffusers Swift 应用

在您自己的苹果硬件上运行 Stable Diffusion 的一个简单方法是使用我们基于 `diffusers` 和苹果的转换与推理仓库的开源 Swift 仓库。您可以研究代码，用 Xcode 编译它，并根据自己的需求进行调整。为方便起见，App Store 中还有一个独立的 Mac 应用，这样您就可以直接体验，而无需处理代码或 IDE。如果您是一名开发者，并确定 Core ML 是构建 Stable Diffusion 应用的最佳解决方案，那么您可以使用本指南的其余部分来开始您的项目。我们迫不及待地想看到您将构建出什么 🙂。