在 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 上執行。我們還展示瞭如何自行轉換微調模型，並迫不及待地想看到您使用這些工具和技術能創造出什麼！