Press "Enter" to skip to content

在Mac上稳定使用高级Core ML量化的Diffusion XL

昨天发布了稳定扩散 XL,它非常棒。它可以生成大尺寸(1024×1024)的高质量图像;通过一些新的技巧改进了对提示的遵循性;由于最新的噪声调度器研究,它可以轻松产生非常暗或非常明亮的图像;而且它是开源的!

不足之处是,该模型更大,因此在消费者硬件上运行速度更慢,更困难。使用最新版本的 Hugging Face 扩散器库,您可以在 16 GB 的 GPU 内存上的 CUDA 硬件上运行稳定扩散 XL,从而可以在 Colab 的免费版本上使用它。

过去几个月表明,人们对于出于各种原因在本地运行 ML 模型非常感兴趣,包括隐私、便利性、更容易的实验或不受计量的使用。我们在 Apple 和 Hugging Face 都在努力探索这个领域。我们已经展示了如何在 Apple Silicon 上运行稳定扩散,或者如何利用 Core ML 的最新进展来改善尺寸和性能,包括 6 位调色板技术。

对于稳定扩散 XL,我们做了几件事情:

  • 将基础模型移植到 Core ML,这样您就可以在您的原生 Swift 应用中使用它。
  • 更新了 Apple 的转换和推断存储库,这样您就可以自己转换模型,包括您感兴趣的任何微调。
  • 更新了 Hugging Face 的演示应用程序,展示如何使用从 Hub 下载的新的 Core ML 稳定扩散 XL 模型。
  • 探索了混合位调色技术,这是一种先进的压缩技术,可以在最小化和控制质量损失的同时实现重要的尺寸减小。您也可以将相同的技术应用于您自己的模型!

所有内容都是开源的,今天就可以使用,让我们开始吧。

目录

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

使用 Hugging Face Hub 中的 SD XL 模型

作为此版本的一部分,我们在 Core ML 中发布了两个不同版本的稳定扩散 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 的演示应用程序来测试任何一个模型。以下是使用新的稳定扩散 XL 流水线的后者的示例:

在Mac上稳定使用高级Core ML量化的Diffusion XL 四海 第1张

与之前的稳定扩散版本一样,我们预计社区会针对不同领域创建新的微调版本,并且其中许多版本将被转换为 Core ML。您可以关注 Hub 中的此筛选器来探索!

稳定扩散 XL 可在运行 macOS 14 的 Apple Silicon Mac 上工作。目前它使用的是适用于 CPU + GPU 计算单元的 ORIGINAL 注意力实现。请注意,尚未移植细化阶段。

作为参考,这是我们在不同设备上实现的性能指标:

什么是混合位调色技术?

上个月我们讨论了 6 位调色技术,这是一种后训练量化方法,将 16 位权重转换为每个参数仅需 6 位。这可以大幅减小模型尺寸,但进一步减小尺寸则更加棘手,因为随着位数的减少,模型质量会受到越来越大的影响。

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

我们尝试了一个不同的替代方案:混合位调色。我们不再每个参数使用6位,而是检查模型并决定每个层使用多少量化位。我们根据每个层对整体质量降低的贡献程度来做出决策,通过比较量化模型和原始模型在float16模式下的峰值信噪比(PSNR)来衡量,对于一组少量输入。我们对每个层进行了几个位深度的探索:1(!),248。如果某个层在使用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。

对于视觉示例,这是在相同种子下运行每个配方时,对于提示一个冲浪狗的高质量照片的结果:

一些初步的结论:

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

这种技术对于Stable Diffusion XL非常有效,因为即使参数数量相对于先前版本增加了三倍,我们仍然可以保持大致相同的UNet大小。但它不仅限于此!您可以将该方法应用于任何Stable Diffusion Core ML模型。

如何创建混合位配方?

下图显示了stabilityai/stable-diffusion-xl-base-1.0的信号强度(以dB表示的PSNR)与模型大小减小(相对于float16大小的百分比)之间的关系。逐渐使用具有固定位数的调色板对更多层进行调色生成了{1,2,4,6,8}位曲线。这些层按照它们对端到端信号强度的独立影响的升序排序,因此累积压缩的影响尽可能地延迟。混合位曲线是基于当一个层的独立影响对端到端信号完整性的降低低于一个阈值时,回退到更高位数的情况。请注意,除了1位外,所有基于调色的曲线在相同的模型大小下都优于线性8位量化。

在Mac上稳定使用高级Core ML量化的Diffusion XL 四海 第2张

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

分析阶段的目标是找到混合位曲线中的点(图中所有其他曲线中的棕色曲线),以便我们可以选择所需的质量与大小的权衡。如前一节所述,我们迭代通过各个层,并选择低于给定PSNR阈值的最低位深度。我们对各种阈值重复此过程,以获得不同的量化策略。该过程的结果是一组量化配方,其中每个配方只是详细说明模型中每个层要使用的位数的JSON字典。为简单起见,忽略参数较少的层,保持为float16。

应用阶段只需按照JSON结构中指定的位数应用调色。

分析是一个耗时的过程,需要GPU(mpscuda),因为我们必须多次运行推断。一旦完成,可以在几分钟内执行配方应用。

我们为每个阶段提供了脚本:

  • mixed_bit_compression_pre_analysis.py
  • mixed_bit_compression_apply.py

转换精调模型

如果您之前已将稳定扩散模型转换为Core ML,则使用命令行转换器将XL模型转换的过程非常相似。有一个新的标志用于指示模型是否属于XL系列,如果是,则必须使用--attention-implementation ORIGINAL

有关该过程的介绍,请查看存储库中的说明或我们以前的博客帖子之一,并确保使用上述标志。

运行混合位调色

将稳定扩散或稳定扩散XL模型转换为Core ML后,您可以选择使用上述脚本应用混合位调色。

由于分析过程较慢,我们已经为最流行的模型准备了配方:

  • 稳定扩散1.5的配方
  • 稳定扩散2.1的配方
  • 稳定扩散XL 1.0基础的配方

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

此外,我们还将稳定扩散XL分析的三个最佳配方应用于UNet的Core ML版本,并在此处发布。请随意尝试并查看它们对您的工作效果如何!

最后,如介绍中所述,我们创建了一个完整的稳定扩散XL Core ML流水线,使用4.5-bit配方。

发布资源

  • apple/ml-stable-diffusion,由Apple提供。用于Swift(和Python)的转换和推断库。
  • huggingface/swift-coreml-diffusers。基于Apple包构建的Hugging Face演示应用程序。
  • 稳定扩散XL 1.0基础(Core ML版本)。使用上述存储库和其他第三方应用程序准备运行的模型。
  • 带有混合位调色的稳定扩散XL 1.0基础(Core ML)。与上述相同的模型,使用4.5位的有效调色量化(平均值)对UNet进行量化。
  • 其他带有混合位调色的UNet。
  • 混合位调色配方,预先计算适用于常见模型并可直接使用。
  • mixed_bit_compression_pre_analysis.py。用于运行混合位分析和配方生成的脚本。
  • mixed_bit_compression_apply.py。用于应用在分析阶段计算的配方的脚本。
Leave a Reply

Your email address will not be published. Required fields are marked *