在 Mac 上使用 Advanced Core ML Quantization 运行 Stable Diffusion XL

Stable Diffusion XL 于昨天发布，它非常棒。它可以生成大型（1024x1024）高质量图像；通过一些新技巧，对提示的遵循度得到了提高；得益于噪声调度器方面的最新研究，它可以毫不费力地生成非常暗或非常亮的图像；而且它是开源的！

缺点是模型更大，因此在消费级硬件上运行速度更慢、更困难。使用Hugging Face diffusers 库的最新版本，您可以在具有 16 GB GPU RAM 的 CUDA 硬件上运行 Stable Diffusion XL，使其可以在 Colab 的免费层上使用。

过去几个月表明，人们出于各种原因（包括隐私、便利、更轻松的实验或不限量使用）对在本地运行机器学习模型非常感兴趣。Apple 和 Hugging Face 都在努力探索这一领域。我们已经展示了如何在 Apple Silicon 上运行 Stable Diffusion，或者如何利用Core ML 的最新进展，通过 6 位调色来提高大小和性能。

对于 Stable Diffusion XL，我们做了一些事情：

• 将基础模型移植到 Core ML，以便您可以在您的原生 Swift 应用程序中使用它。
• 更新了Apple 的转换和推理仓库，以便您可以自己转换模型，包括您感兴趣的任何微调模型。
• 更新了Hugging Face 的演示应用程序，以展示如何使用从 Hub 下载的新 Core ML Stable Diffusion XL 模型。
• 探索了混合位调色，这是一种先进的压缩技术，可以在最大限度地减少和控制质量损失的情况下实现重要的尺寸减小。您也可以将相同的技术应用于您自己的模型！

一切都是开源的，今天即可使用，让我们开始吧。

目录

• 使用 Hugging Face Hub 中的 SD XL 模型
• 什么是混合位调色？
• 混合位配方是如何创建的？
• 转换微调模型
• 已发布资源

使用 Hugging Face Hub 中的 SD XL 模型

作为此次发布的一部分，我们在 Core ML 中发布了两个不同版本的 Stable Diffusion XL。

• apple/coreml-stable-diffusion-xl-base 是一个完整的管道，没有任何量化。
• apple/coreml-stable-diffusion-mixed-bit-palettization 包含（除其他工件外）一个完整的管道，其中 UNet 已被混合位调色“配方”取代，该配方实现了相当于每参数 4.5 位的压缩。大小从 4.8 GB 降至 1.4 GB，减少了 71%，我们认为质量仍然很棒。

这两种模型都可以使用 Apple 的Swift 命令行推理应用程序或 Hugging Face 的演示应用程序进行测试。下面是使用新的 Stable Diffusion XL 管道的后者的一个示例：

与之前的 Stable Diffusion 版本一样，我们预计社区将为不同领域开发新的微调版本，其中许多将转换为 Core ML。您可以关注Hub 中的此过滤器进行探索！

Stable Diffusion XL 可以在运行 macOS 14 公测版的 Apple Silicon Mac 上运行。它目前使用 `ORIGINAL` 注意力实现，该实现适用于 CPU + GPU 计算单元。请注意，refiner 阶段尚未移植。

以下是我们不同设备上实现的性能数据：

什么是混合位调色？

上个月我们讨论了 6 位调色，这是一种训练后量化方法，将 16 位权重转换为每参数仅 6 位。这实现了模型尺寸的显著减小，但进一步缩小则很棘手，因为随着位数减少，模型质量受到的影响越来越大。

进一步减小模型尺寸的一个选择是使用训练时量化，它包括在微调模型时学习量化表。这很有效，但您需要为您想要转换的每个模型运行一个微调阶段。

我们探索了另一种替代方案：混合位调色。我们不是对每个参数使用 6 位，而是检查模型并决定每个层使用多少量化位。我们根据每个层对整体质量下降的贡献程度做出决定，我们通过比较量化模型和原始模型在 `float16` 模式下对一组少量输入的 PSNR 来衡量。我们探索每个层的几种位深度：`1` (!) 、`2`、`4` 和 `8`。如果一个层在使用（比如说）2 位时显著退化，我们就移动到 `4`，依此类推。如果某些层对保持质量至关重要，它们可能会保留在 16 位模式下。

使用这种方法，我们可以实现平均 2.8 位等有效量化，并且我们测量了我们尝试的每种组合对退化的影响。这使我们能够更好地了解用于目标质量和尺寸预算的最佳量化。

为了说明该方法，让我们考虑一下我们从一次分析运行中获得的以下量化“配方”（我们稍后将解释它们是如何生成的）

{
  "model_version": "stabilityai/stable-diffusion-xl-base-1.0",
  "baselines": {
    "original": 82.2,
    "linear_8bit": 66.025,
    "recipe_6.55_bit_mixedpalette": 79.9,
    "recipe_4.50_bit_mixedpalette": 75.8,
    "recipe_3.41_bit_mixedpalette": 71.7,
  },
}

这告诉我们，在 float16 中测量的原始模型质量约为 82 dB。执行朴素的 8 位线性量化会将其降至 66 dB。但我们有一个配方，平均每参数压缩到 6.55 位，同时保持 PSNR 为 80 dB。第二个和第三个配方进一步减小了模型大小，同时仍保持 PSNR 大于 8 位线性量化。

对于视觉示例，这是在相同种子下运行三个配方中每个配方的提示“一只冲浪狗的高质量照片”的结果

一些初步结论

• 我们认为，所有图像在逼真度方面都具有良好的质量。6.55 位和 4.50 位版本在这方面接近 16 位版本。
• 相同的种子会产生等效的构图，但不会保留相同的细节。例如，狗的品种可能不同。
• 随着我们增加压缩，对提示的遵循度可能会下降。在此示例中，激进的 3.41 位版本丢失了板。PSNR 仅比较像素的整体差异程度，而不关心图像中的主体。您需要检查结果并根据您的用例进行评估。

这种技术对于 Stable Diffusion XL 非常有用，因为即使参数数量相对于之前的版本增加了三倍，我们也可以保持大致相同的 UNet 大小。但它并非独有！您可以将此方法应用于任何 Stable Diffusion Core ML 模型。

混合位配方是如何创建的？

下图显示了 `stabilityai/stable-diffusion-xl-base-1.0` 的信号强度（PSNR，单位 dB）与模型尺寸减小（float16 尺寸的百分比）的关系。`{1,2,4,6,8}` 位曲线是通过使用固定位数调色板逐步调色更多层生成的。层按其对端到端信号强度的独立影响的升序排列，因此累积压缩的影响尽可能延迟。混合位曲线基于一旦层的独立影响降至阈值以下，就回退到更高的位数。请注意，除 1 位外，所有基于调色的曲线在相同模型尺寸下均优于线性 8 位量化。

混合位调色分两个阶段运行：分析和应用。

分析阶段的目标是找到混合位曲线中的点（图中高于所有其他曲线的棕色曲线），以便我们选择所需的质量与尺寸权衡。如上一节所述，我们迭代各层并选择在给定 PSNR 阈值之上产生结果的最低位深度。我们对各种阈值重复该过程以获得不同的量化策略。因此，该过程的结果是一组量化配方，其中每个配方都只是一个 JSON 字典，详细说明了模型中每个层使用的位数。参数较少的层为简单起见而被忽略并保留在 float16 中。

应用阶段只是遍历配方并使用 JSON 结构中指定的位数进行调色。

分析是一个耗时的过程，需要 GPU（`mps` 或 `cuda`），因为我们必须多次运行推理。一旦完成，配方应用可以在几分钟内完成。

我们为这些阶段的每个阶段都提供了脚本：

• mixed_bit_compression_pre_analysis.py
• mixed_bit_compression_apply.py

转换微调模型

如果您以前将 Stable Diffusion 模型转换为 Core ML，那么使用命令行转换器转换 XL 的过程非常相似。有一个新标志可以指示模型是否属于 XL 系列，如果属于，则必须使用 `—attention-implementation ORIGINAL`。

有关该过程的介绍，请查看仓库中的说明或我们之前的一篇博客文章，并确保使用上述标志。

运行混合位调色

将 Stable Diffusion 或 Stable Diffusion XL 模型转换为 Core ML 后，您可以选择使用上述脚本应用混合位调色。

由于分析过程缓慢，我们为最流行的模型准备了配方：

• Stable Diffusion 1.5 的配方
• Stable Diffusion 2.1 的配方
• Stable Diffusion XL 1.0 基础版的配方

您可以下载并在本地应用它们进行实验。

此外，我们还将 Stable Diffusion XL 分析中的三个最佳配方应用于 UNet 的 Core ML 版本，并将其发布在此处。欢迎使用它们，看看它们对您有何帮助！

最后，如引言中所述，我们创建了一个使用 `4.5 位` 配方的完整 Stable Diffusion XL Core ML 管道。

已发布资源

• apple/ml-stable-diffusion，由 Apple 提供。Swift（和 Python）的转换和推理库。
• huggingface/swift-coreml-diffusers。Hugging Face 演示应用程序，基于 Apple 的包构建。
• Stable Diffusion XL 1.0 基础版（Core ML 版本）。模型已准备好使用上述仓库和其他第三方应用程序运行。
• Stable Diffusion XL 1.0 基础版，带混合位调色（Core ML）。与上述模型相同，其 UNet 经过量化，有效调色为 4.5 位（平均）。
• 额外的带有混合位调色的 UNet 模型.
• 混合位调色配方，为流行模型预计算并准备使用。
• mixed_bit_compression_pre_analysis.py。用于运行混合位分析和配方生成的脚本。
• mixed_bit_compression_apply.py。用于应用在分析阶段计算的配方的脚本。