在 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。用於應用在分析階段計算的配方的指令碼。