如何在 iPhone、iPad 和 Mac 上通过 Core ML 加速 Stable Diffusion

WWDC’23（Apple 全球开发者大会）于上周举行。很多新闻都聚焦于主题演讲中 Vision Pro 的发布，但内容远不止于此。和每年一样，WWDC 周充满了 200 多场技术会议，深入探讨 Apple 操作系统和框架即将推出的功能。今年我们尤其对 Core ML 中专注于压缩和优化技术的改动感到兴奋。这些改动让运行诸如 Stable Diffusion 之类的模型变得更快，并且占用更少内存！作为一点体验，不妨看看我去年 12 月在我的 iPhone 13 上运行的以下测试，并与当前使用 6 位调色板量化后的速度进行比较：

iPhone 上的 Stable Diffusion，去年 12 月与现在使用 6 位调色板量化的对比

目录

• 新的 Core ML 优化
• 使用量化和优化后的 Stable Diffusion 模型
• 转换和优化自定义模型
• 使用低于 6 位
• 结论

新的 Core ML 优化

Core ML 是一个成熟的框架，允许机器学习模型在设备端高效运行，并充分利用 Apple 设备中的所有计算硬件：CPU、GPU 以及专门用于 ML 任务的神经引擎。由于 Stable Diffusion 和带有聊天界面的大型语言模型等模型的流行，设备端执行正经历一段异常受关注的时期。许多人出于各种原因希望在自己的硬件上运行这些模型，包括便利性、隐私以及节省 API 成本。自然地，许多开发者正在探索在设备端高效运行这些模型的方法，并创造新的应用和用例。有助于实现这一目标的 Core ML 改进对整个社区来说都是重大新闻！

Core ML 的优化改动涵盖两个不同（但互补）的软件包：

• Core ML 框架本身。这是在 Apple 硬件上运行 ML 模型的引擎，也是操作系统的一部分。模型必须以该框架支持的特殊格式导出，这种格式也被称为“Core ML”。
• coremltools 转换包。这是一个开源 Python 模块，其使命是将 PyTorch 或 Tensorflow 模型转换为 Core ML 格式。

coremltools 现在包含一个名为 coremltools.optimize 的新子模块，其中包含所有压缩和优化工具。有关这个软件包的完整细节，请查看这场 WWDC 会议。就 Stable Diffusion 而言，我们将使用6 位调色板量化，这是一种量化方式，可将模型权重从 16 位浮点表示压缩到每个参数仅 6 位。“调色板量化”这个名称指的是一种类似于计算机图形学中用于处理有限颜色集的技术：颜色表（或“调色板”）包含固定数量的颜色，图像中的颜色会被替换为调色板中可用最接近颜色的索引。这立即带来了大幅减小存储大小的好处，从而减少下载时间和设备端磁盘使用。

2 位调色板量化示意图。图片来源：Apple WWDC’23 会议 Use Core ML Tools for machine learning model compression。

压缩后的 6 位权重无法用于计算，因为它们只是表中的索引，不再代表原始权重的大小。因此，Core ML 在使用前需要先解压调色板化的权重。在之前版本的 Core ML 中，解压发生在模型首次从磁盘加载时，所以占用的内存量等于解压后的模型大小。有了新的改进，权重会保持为 6 位数字，并在推理逐层进行时即时转换。这看起来可能很慢——一次推理运行需要大量解压操作——但通常比以 16 位模式准备所有重量更高效！原因是内存传输处于执行的关键路径上，而传输更少的内存比传输未压缩数据更快。

使用量化和优化的 Stable Diffusion 模型

去年 12 月，Apple 推出了 ml-stable-diffusion，这是一个基于 diffusers 的开源仓库，可轻松将 Stable Diffusion 模型转换为 Core ML。它还对 transformer 注意力层应用了优化，从而在 Neural Engine（在可用设备上）上加快推理速度。ml-stable-diffusion 已在 WWDC 后更新，内容如下：

• 在转换过程中支持使用 --quantize-nbits 进行量化。你可以量化到 8、6、4，甚至 2 位！为了获得最佳效果，我们建议使用 6 位量化，因为精度损失很小，同时还能实现快速推理并显著节省内存。如果你想低于这个位数，请查看本节了解高级技术。
• 对注意力层进行了额外优化，在 Neural Engine 上实现了更好的性能！技巧是将查询序列拆分为 512 的块，以避免创建大型中间张量。这种方法在代码中称为 SPLIT_EINSUM_V2，可以将性能提高 10% 到 30%。

为了让每个人都能轻松利用这些改进，我们转换了四个官方 Stable Diffusion 模型，并将它们推送到 Hub。以下是所有变体：

为了使用 6 位模型，你需要 iOS/iPadOS 17 或 macOS 14（Sonoma）的开发版本，因为这些版本包含最新的 Core ML 框架。如果你是注册开发者，可以从 Apple 开发者网站下载它们，或者你可以注册几周后将发布的公开测试版。

请注意，每个变体都提供 Core ML 格式，也提供 zip 归档。Zip 文件非常适合原生应用，例如我们的开源演示应用和其他第三方工具。如果你只想在自己的硬件上运行模型，最简单的方法是使用我们的演示应用，并选择你想测试的量化模型。你需要使用 Xcode 编译该应用，但很快 App Store 中将提供可下载的更新。更多详情，请查看我们之前的文章。

在演示应用中运行 6 位 stable-diffusion-2-1-base 模型

如果你想下载特定的 Core ML 包以集成到你自己的 Xcode 项目中，你可以克隆这些仓库，或者使用如下代码下载你感兴趣的版本。

from huggingface_hub import snapshot_download
from pathlib import Path

repo_id = "apple/coreml-stable-diffusion-2-1-base-palettized"
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}")

转换和优化自定义模型

如果你想使用个性化的 Stable Diffusion 模型（例如，如果你已经微调或对模型进行了 dreambooth 化），你可以使用 Apple 的 ml-stable-diffusion 仓库自行进行转换。以下是操作方法的简要概述，但我们建议你阅读文档详情。

如果你想应用量化，你需要最新版本的 coremltools、apple/ml-stable-diffusion 和 Xcode 才能进行转换。

• 从 GitHub 的发布页面下载 coremltools 7.0 beta。
• 从Apple 开发者网站下载 Xcode 15.0 beta。
• 从仓库下载 apple/ml-stable-diffusion 并按照安装说明进行操作。

1. 选择你想要转换的模型。你可以训练自己的模型，或从 HF Mirror Diffusers 模型库中选择一个。例如，让我们转换 prompthero/openjourney-v4。
2. 安装 apple/ml-stable-diffusion 并使用 ORIGINAL 注意力实现运行首次转换，如下所示：

python -m python_coreml_stable_diffusion.torch2coreml \
    --model-version prompthero/openjourney-v4 \
    --convert-unet \
    --convert-text-encoder \
    --convert-vae-decoder \
    --convert-vae-encoder \
    --convert-safety-checker \
    --quantize-nbits 6 \
    --attention-implementation ORIGINAL \
    --compute-unit CPU_AND_GPU \
    --bundle-resources-for-swift-cli \
    --check-output-correctness \
    -o models/original/openjourney-6-bit

• 如果你想使用图像到图像任务，请使用 --convert-vae-encoder。
• 不要将 --chunk-unet 与 --quantized-nbits 6（或更低版本）一起使用，因为量化后的模型足够小，在 iOS 和 macOS 上都能很好地工作。

1. 针对 SPLIT_EINSUM_V2 注意力实现重复转换：

python -m python_coreml_stable_diffusion.torch2coreml \
    --model-version prompthero/openjourney-v4 \
    --convert-unet \
    --convert-text-encoder \
    --convert-vae-decoder \
    --convert-safety-checker \
    --quantize-nbits 6 \
    --attention-implementation SPLIT_EINSUM_V2 \
    --compute-unit ALL \
    --bundle-resources-for-swift-cli \
    --check-output-correctness \
    -o models/split_einsum_v2/openjourney-6-bit

1. 在目标硬件上测试转换后的模型。根据经验，ORIGINAL 版本在 macOS 上通常表现更好，而 SPLIT_EINSUM_V2 在 iOS 上通常更快。有关更多细节和额外数据点，请参阅社区贡献的这些测试，它们针对的是之前版本的适用于 Core ML 的 Stable Diffusion。

2. 要将所需模型集成到你自己的应用中：

   • 如果你打算在应用内分发模型，请使用 .mlpackage 文件。请注意，这会增加应用二进制文件的大小。
   • 否则，你可以使用编译后的 Resources 在应用启动时动态下载它们。

如果你不使用 --quantize-nbits 选项，权重将表示为 16 位浮点数。这与当前版本的 Core ML 兼容，因此你无需安装 iOS、macOS 或 Xcode 的 beta 版。

使用低于 6 位

6 位量化是模型质量、模型大小和便利性之间的最佳平衡点——你只需提供一个转换选项，就能对任何预训练模型进行量化。这是一个训练后压缩的示例。

上周发布的 coremltools beta 版还包含训练时压缩方法。这里的思路是，你可以微调预训练的 Stable Diffusion 模型，并在微调进行时执行权重压缩。这使你能够使用 4 位甚至 2 位压缩，同时将质量损失降至最低。之所以有效，是因为权重聚类是使用可微算法执行的，因此我们可以应用常规的训练优化器来找到量化表，同时最小化模型损失。

我们计划很快评估这种方法，并且迫不及待地想看到 4 位优化模型的效果以及它们的运行速度。如果你抢先做到了，请给我们留言，我们会很乐意查看 🙂

结论

量化方法可用于缩小 Stable Diffusion 模型的体积，使其在设备端运行更快、消耗资源更少。最新版本的 Core ML 和 coremltools 支持 6 位调色板化等技术，这些技术易于应用，且对质量影响极小。我们已将 6 位调色板化模型添加到 Hub，其体积小到足以在 iOS 和 macOS 上运行。我们还展示了如何自行转换微调后的模型，并非常期待看到你们用这些工具和技术做出什么成果！