Hugging Face 讲解如何用 Core ML 在 iPhone、iPad 和 Mac 上更快运行 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 以及专为机器学习任务设计的神经引擎。设备端执行正经历一段异常火热的时期，这由 Stable Diffusion 和带聊天界面的大语言模型等模型的流行所引发。出于便利性、隐私和 API 成本节省等各种原因，许多人都希望在自己的硬件上运行这些模型。自然而然地，许多开发者正在探索在设备端高效运行这些模型的方法，并创造新的应用和用例。有助于实现这一目标的 Core ML 改进对整个社区来说都是重大新闻！

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

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

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

2 位调色板量化示意图。图片来源：Apple WWDC’23 讲座 使用 Core ML Tools 进行机器学习模型压缩。

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

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

去年 12 月，Apple 推出了 ml-stable-diffusion，这是一个基于 diffusers 的开源仓库，可轻松将 Stable Diffusion 模型转换为 Core ML。它还应用了对 transformers 注意力层的优化，使推理在 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. 选择你想要转换的模型。你可以训练自己的模型，或者从Hugging Face 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 上通常更快。有关更多细节和额外数据点，请参阅社区贡献的这些测试，它们针对的是之前版本的 Stable Diffusion for Core ML。

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 上运行。我们还展示了如何自行转换微调后的模型，并迫不及待想看到你们用这些工具和技术做出什么！