Stable Diffusion XL 登陆 Mac：Core ML 混合位量化实现 71% 体积压缩

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

缺点是模型更大，因此在消费级硬件上运行更慢、更困难。使用 HF Mirror diffusers 库的最新版本，你可以在 16 GB GPU 内存的 CUDA 硬件上运行 Stable Diffusion XL，从而可以在 Colab 的免费套餐上使用它。

过去几个月表明，人们出于各种原因——包括隐私、便利、更轻松的实验或无计量使用——非常明显地对在本地运行 ML 模型感兴趣。我们一直在 Apple 和 HF Mirror 努力工作，探索这一领域。我们已经展示了 如何在 Apple Silicon 上运行 Stable Diffusion，或者如何利用 Core ML 的最新进展，通过 6 位调色板化来改进大小和性能。

对于 Stable Diffusion XL，我们做了几件事：

• 将基础模型移植到 Core ML，以便你可以在原生 Swift 应用中使用它。
• 更新了 Apple 的转换和推理仓库，以便你可以自行转换模型，包括你感兴趣的任何微调版本。
• 更新了 HF Mirror 的演示应用，以展示如何使用从 Hub 下载的新 Core ML Stable Diffusion XL 模型。
• 探索了混合位调色板化，这是一种先进的压缩技术，可以在实现重要的大小缩减的同时，最小化并控制你遭受的质量损失。你也可以将同样的技术应用于你自己的模型！

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

目录

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

使用 HF Mirror 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 命令行推理应用或 HF Mirror 的演示应用进行测试。以下是后者使用新 Stable Diffusion XL 管道的示例：

与之前的 Stable Diffusion 版本一样，我们期望社区为不同领域提出新颖的微调版本，其中许多将被转换为 Core ML。你可以关注Hub 中的这个过滤器来探索！

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

作为参考，这些是我们在不同设备上达到的性能数据：

什么是混合位调色板化？

上个月我们讨论了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中的PSNR衡量）约为82 dB。执行简单的8位线性量化会将其降至66 dB。但接着我们有一个配方，平均每个参数压缩到6.55位，同时将PSNR保持在80 dB。第二个和第三个配方进一步减小了模型大小，同时仍保持了高于8位线性量化的PSNR。

对于视觉示例，这些是在提示a high quality photo of a surfing dog上使用相同种子运行三个配方中的每一个的结果：

3.41位	4.50位	6.55位	16位（原始）

一些初步结论：

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

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

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

下图展示了 stabilityai/stable-diffusion-xl-base-1.0 的模型大小缩减（float16 大小的百分比）与信号强度（PSNR，以 dB 为单位）之间的关系。{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 base 的配方

您可以下载这些配方并在本地应用以进行实验。

此外，我们还将 Stable Diffusion XL 分析中最好的三个配方应用于 Core ML 版本的 UNet，并在此处发布。欢迎试用，看看它们对您效果如何！

最后，如引言所述，我们创建了一个完整的 Stable Diffusion XL Core ML 流水线，它使用 4.5-bit 配方。

已发布资源

• apple/ml-stable-diffusion，由 Apple 提供。用于 Swift（和 Python）的转换和推理库。
• huggingface/swift-coreml-diffusers。HF Mirror 演示应用，基于 Apple 的软件包构建。
• Stable Diffusion XL 1.0 base（Core ML 版本）。使用上述仓库及其他第三方应用即可运行的模型。
• Stable Diffusion XL 1.0 base，采用混合位调色板量化（Core ML）。与上述模型相同，其中 UNet 采用了平均有效调色板为 4.5 位的量化。
• 附加的采用混合位调色板量化的 UNet。
• 混合位调色板量化配方，为热门模型预先计算，开箱即用。
• mixed_bit_compression_pre_analysis.py。用于运行混合位分析与配方生成的脚本。
• mixed_bit_compression_apply.py。用于应用分析阶段计算出的配方的脚本。