在 iPhone、iPad 和 Mac 上使用 Core ML 加速 Stable Diffusion

WWDC '23（Apple 全球开发者大会）已于上周举行。许多新闻都集中在主题演讲中 Vision Pro 的发布上，但其内容远不止于此。像每年一样，WWDC 周包含 200 多个技术会议，深入探讨 Apple 操作系统和框架的未来功能。今年，我们特别兴奋于 Core ML 在压缩和优化技术方面的变化。这些变化使得运行 Stable Diffusion 等模型更快，并使用更少的内存！作为示例，请看我在去年 12 月在我的 iPhone 13 上运行的以下测试，与目前使用 6 位调色技术的速度进行比较。

Stable Diffusion 在 iPhone 上的表现，去年 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 会议 使用 Core ML 工具进行机器学习模型压缩。

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

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

去年 12 月，苹果推出了基于 diffusers 的开源仓库 ml-stable-diffusion，用于轻松将 Stable Diffusion 模型转换为 Core ML。它还对transformers 注意力层应用了优化，从而在神经引擎（在可用设备上）上实现更快的推理。WWDC 后，ml-stable-diffusion 刚刚更新了以下内容：

• 转换期间使用 --quantize-nbits 支持量化。您可以量化为 8、6、4 甚至 2 位！为了获得最佳结果，我们建议使用 6 位量化，因为精度损失很小，同时实现了快速推理和显著的内存节省。如果您想低于此值，请查看本节以获取高级技术。
• 注意力层的额外优化，在神经引擎上实现更好的性能！诀窍是将查询序列分成 512 的块，以避免创建大型中间张量。此方法在代码中称为 SPLIT_EINSUM_V2，可以将性能提高 10% 到 30%。

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

请注意，每个变体都提供 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 模型（例如，如果您已微调或 dreamboothed 自己的模型），您可以使用 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

3. 为 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

4. 在所需的硬件上测试转换后的模型。经验法则，ORIGINAL 版本通常在 macOS 上表现更好，而 SPLIT_EINSUM_V2 通常在 iOS 上更快。有关更多详细信息和附加数据点，请参阅社区在 Core ML 早期版本的 Stable Diffusion 上贡献的这些测试。

5. 将所需模型集成到您自己的应用程序中

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

使用少于 6 位

6 位量化在模型质量、模型大小和便利性之间找到了一个最佳平衡点——您只需提供一个转换选项即可量化任何预训练模型。这是一个*训练后压缩*的例子。

上周发布的 coremltools 测试版还包括*训练时*压缩方法。其思想是，您可以在微调预训练的 Stable Diffusion 模型时进行权重压缩。这允许您使用 4 位甚至 2 位压缩，同时最大限度地减少质量损失。之所以可行，是因为权重聚类是使用可微分算法执行的，因此我们可以应用常规的训练优化器来找到量化表，同时最小化模型损失。

我们计划很快评估此方法，并迫不及待地想看看 4 位优化模型的效果和运行速度。如果您在这方面领先于我们，请给我们留言，我们很乐意查看 🙂

结论

量化方法可用于减小 Stable Diffusion 模型的大小，使其在设备上运行更快并消耗更少的资源。Core ML 和 coremltools 的最新版本支持 6 位调色等技术，这些技术易于应用且对质量影响最小。我们已将 6 位调色模型添加到 Hub，这些模型足够小，可以在 iOS 和 macOS 上运行。我们还展示了如何自行转换微调模型，并迫不及待地想看到您使用这些工具和技术能创造出什么！